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

from collection import Collection, Scorer
from random import randint
from clustering import Partition
from distance_matrix import DistanceMatrix
from lib.remote.externals.phyml import Phyml

class emtrees(object):

    # Not sure if scorer required
    # Possible end criteria?

    def __init__(
        self,
        collection,
        nclusters,
        metric = 'euc',
        method='distance',
        ):

        if not isinstance(nclusters, int) or nclusters <= 1:
            raise Exception('Need appropriate value for number of clusters.')

        self.nclusters = nclusters
        self.scorer = Scorer(collection.records, collection.analysis) # Could check for entries
        self.datatype = collection.datatype
        self.tmpdir = collection.tmpdir
        self.method = method
        self.metric = metric
        self.assign_clusters()

    def assign_clusters(self):
        # Collapse equivalent trees?
        k = self.nclusters
        self.partition = Partition( [randint(1,k) for rec in self.scorer.records] )
        self.L = self.scorer.score(self.partition)

    def maximise(self): # Needs to change
        count = 0
        # clusters = [ Cluster(self.partition.get_membership()[n], self.scorer,records, self.scorer.analysis) for n in range(self.clusters)]
        clusters = [0] * self.nclusters

        while True:

            for n in range(self.nclusters):
                members = self.partition.get_membership()[n]
                if not clusters[n] or clusters[n].members != members:
                    clusters[n] = Cluster(members, self.scorer.records, self.scorer.analysis)

            assignment = getattr(self,self.method)(clusters)
            score = self.scorer.score(assignment)

            if score > self.L:
                self.L = score
                self.partition = assignment

            else:
                count += 1
                if count > 1: # Algorithm is deterministic so no need for more iterations
                    break

    def distance(self,clusters): # need better name once we figure out what it does

        assignment = self.partition.partition_vector[:]

        for (index, record) in enumerate(self.scorer.records):
            dists = [ self.dist(record.tree, clusters[n].tree) for n in range(self.nclusters) ]
            if assignment.count(assignment[index]) > 1:
                assignment[index] = dists.index(min(dists)) + 1

        return(Partition(assignment))

    def ml(self,clusters):

        assignment = self.partition.partition_vector[:]

        for (index, record) in enumerate(self.scorer.records):
            likelioods = [ self.phyml_likelihood(record, clusters[n]) for n in range(self.nclusters) ]
            if assignment.count(assignment[index]) > 1:
                assignment[index] = likelioods.index(min(likelioods)) + 1

        return(Partition(assignment))

    def phyml_likelihood(self, record, cluster, verbose=0):
        p = Phyml(record, tmpdir=self.tmpdir)
        cluster.tree.write_to_file('test_tree')
        self.add_tempfile('test_tree')
        p.add_flag('--inputtree', 'test_tree') # Need tempdir????
        p.add_flag('-o', 'r') # Optimise only banch lengtha and substitutions`
        p.add_flag('-a', 'e')
        p.add_flag('-b', 0)
        p.add_flag('-c', 4)
        p.add_flag('-d', 'nt') # set data type - could inherit from Collection object?
        p.run(verbosity=verbose)

    def dist(self, tree1, tree2):
        distance = DistanceMatrix( [tree1, tree2], self.metric)[0][1]
        return(distance)

class Cluster(object):
    def __init__(self, members, records, analysis):
        self.members = tuple(members)
        self.records = [ records[i] for i in self.members ]
        self.scorer = Scorer(records, analysis)
        self.tree = self.scorer.add(self.members)
<commit_msg>added working ml methods<commit_after>#!/usr/bin/env python

from collection import Collection, Scorer
from random import randint
from clustering import Partition
from distance_matrix import DistanceMatrix
from lib.remote.externals.phyml import Phyml

class emtrees(object):

    # Not sure if scorer required
    # Possible end criteria?

    def __init__(
        self,
        collection,
        nclusters,
        metric = 'euc',
        ):

        if not isinstance(nclusters, int) or nclusters <= 1:
            raise Exception('Need appropriate value for number of clusters.')

        self.nclusters = nclusters
        self.scorer = Scorer(collection.records, collection.analysis) # Could check for entries
        self.datatype = collection.datatype
        self.tmpdir = collection.tmpdir
        self.metric = metric

    def assign_partition(self):
        # Collapse equivalent trees?
        k = self.nclusters
        self.partition = Partition( [randint(1,k) for rec in self.scorer.records] )
        self.L = self.scorer.score(self.partition)

    def assign_clusters(self,clusters,partition):
        for n in range(self.nclusters):
            members = partition.get_membership()[n]
            if not clusters[n] or clusters[n].members != members:
                clusters[n] = Cluster(members, self.scorer.records, self.scorer.analysis)

        return(clusters)

    def maximise(self, method): # Needs to change
        self.assign_partition()
        clusters = [0] * self.nclusters
        alg = getattr(self,method)
        count = 0
        # clusters = [ Cluster(self.partition.get_membership()[n], self.scorer,records, self.scorer.analysis) for n in range(self.clusters)]

        while True:
            self.assign_clusters(clusters,self.partition)
            assignment = list(self.partition.partition_vector)

            for (index, record) in enumerate(self.scorer.records):
                scores = [ alg(record, clusters[n]) for n in range(self.nclusters) ]
                print scores
                if assignment.count(assignment[index]) > 1:
                    assignment[index] = scores.index(max(scores)) + 1

            assignment = Partition(assignment)
            score = self.scorer.score(assignment)

            if score > self.L:
                self.L = score
                self.partition = assignment

            else:
                count += 1
                if count > 1: break # Algorithm is deterministic so no need for more iterations

    def dist(self, obj1, obj2):
        distance = DistanceMatrix( [obj1.tree, obj2.tree], self.metric)[0][1]
        return(-distance)

    def ml(self, record, cluster, verbose=1):
        p = Phyml(record, tmpdir=self.tmpdir)
        cluster.tree.write_to_file('test_tree')
        p.add_tempfile('test_tree')
        p.add_flag('--inputtree', 'test_tree') # Need tempdir????
        p.add_flag('-o', 'r') # Optimise only on substitutions`
        p.add_flag('-a', 'e')
        p.add_flag('-b', 0)
        p.add_flag('-c', 4)
        if self.datatype == 'protein':
            p.add_flag('-d', 'aa') # set data type - could inherit from Collection object?
        elif self.datatype == 'dna':
            p.add_flag('-d', 'nt')
        tree = p.run(verbosity=verbose)
        return(tree.score)

class Cluster(object):
    def __init__(self, members, records, analysis):
        self.members = tuple(members)
        self.records = [ records[i] for i in self.members ]
        self.scorer = Scorer(records, analysis)
        self.tree = self.scorer.add(self.members)
<|endoftext|>


<commit_before>#  Copyright 2008-2012 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import os

from .markuputils import html_escape, xml_escape, attribute_escape

class _MarkupWriter(object):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        if isinstance(output, basestring):
            output = open(output, 'w')
        self.output = output
        self._line_separator = line_separator
        self._encode = self._create_encoder(encoding)

    def _create_encoder(self, encoding):
        if encoding is None:
            return lambda text: text
        return lambda text: text.encode(encoding)

    def start(self, name, attrs=None, newline=True):
        self._write('<%s%s>%s' % (name, self._get_attrs(attrs),
                                  self._line_separator if newline else ''))

    def content(self, content=None, escape=True):
        if content:
            self._write(self._escape(content) if escape else content)

    def end(self, name, newline=True):
        self._write('</%s>%s' % (name, self._line_separator if newline else ''))

    def element(self, name, content=None, attrs=None, escape=True,
                newline=True):
        self.start(name, attrs, newline=False)
        self.content(content, escape)
        self.end(name, newline)

    def close(self):
        self.output.close()

    def _write(self, text):
        self.output.write(self._encode(text))

    def _get_attrs(self, attrs):
        if not attrs:
            return ''
        return ' ' + ' '.join(self._format_attributes(attrs))

class HtmlWriter(_MarkupWriter):

    def _escape(self, content):
        return html_escape(content)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(attrs[name]))
                             for name in sorted(attrs))

class XmlWriter(_MarkupWriter):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        _MarkupWriter.__init__(self, output, line_separator, encoding)
        self._preamble()

    def _preamble(self):
        self.content('<?xml version="1.0" encoding="UTF-8"?>\n', escape=False)

    def _escape(self, text):
        return xml_escape(text)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(unicode(attrs[name])))
                             for name in attrs)
<commit_msg>markupwriter: api documentation<commit_after>#  Copyright 2008-2012 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import os

from .markuputils import html_escape, xml_escape, attribute_escape

class _MarkupWriter(object):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        """Creates new _MarkupWriter.

        :param output: Either an opened, file like object, or a path to the
            desired output file. In the latter case, the file is created.
        :param line_separator: Defines the used line separator.
        :param encoding: Encoding to be used to encode all text written to the
            output file. If `None`, text will not be encoded.
        """
        if isinstance(output, basestring):
            output = open(output, 'w')
        self.output = output
        self._line_separator = line_separator
        self._encode = self._create_encoder(encoding)

    def _create_encoder(self, encoding):
        if encoding is None:
            return lambda text: text
        return lambda text: text.encode(encoding)

    def start(self, name, attrs=None, newline=True):
        self._write('<%s%s>%s' % (name, self._get_attrs(attrs),
                                  self._line_separator if newline else ''))

    def content(self, content=None, escape=True):
        if content:
            self._write(self._escape(content) if escape else content)

    def end(self, name, newline=True):
        self._write('</%s>%s' % (name, self._line_separator if newline else ''))

    def element(self, name, content=None, attrs=None, escape=True,
                newline=True):
        self.start(name, attrs, newline=False)
        self.content(content, escape)
        self.end(name, newline)

    def close(self):
        self.output.close()

    def _write(self, text):
        self.output.write(self._encode(text))

    def _get_attrs(self, attrs):
        if not attrs:
            return ''
        return ' ' + ' '.join(self._format_attributes(attrs))

class HtmlWriter(_MarkupWriter):

    def _escape(self, content):
        return html_escape(content)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(attrs[name]))
                             for name in sorted(attrs))

class XmlWriter(_MarkupWriter):

    def __init__(self, output, line_separator=os.linesep, encoding=None):
        _MarkupWriter.__init__(self, output, line_separator, encoding)
        self._preamble()

    def _preamble(self):
        self.content('<?xml version="1.0" encoding="UTF-8"?>\n', escape=False)

    def _escape(self, text):
        return xml_escape(text)

    def _format_attributes(self, attrs):
        return ('%s="%s"' % (name, attribute_escape(unicode(attrs[name])))
                             for name in attrs)
<|endoftext|>


<commit_before>import operator
import sys

import error

class Minimax():
    def __init__(self, board, colors):
        self.board = board
        self.colors = colors
        self.max_player = None

    def _other_color(self, color):
        colors = list(self.colors)
        colors.remove(color)
        return colors[0]

    def best_move(self, color, difficulty):
        player = color
        self.max_player = player
        opponent = self._other_color(color)
        depth = difficulty

        legal_moves = {}
        for column in range(self.board.width):
            if (self.board.valid_move(column, self.board.get_board())
                is not None):
                state = self.make_fake_move(self.board.get_board(),
                                            column, player)
                legal_moves[column] = -self.search(state, depth-1, opponent)
        return max(legal_moves.items(), key=operator.itemgetter(1))

    def make_fake_move(self, state, column, color):
        row = self.board.valid_move(column, state)
        if row is None:
            raise error.InvalidMoveError
        state[column][row] = color
        return state

    def search(self, state, depth, player,
               alpha=-sys.maxsize, beta=sys.maxsize):
        opp = self._other_color(player)
        legal_moves = []
        for column in range(self.board.width):
            if self.board.valid_move(column, state) is not None:
                legal_moves.append(self.make_fake_move(state, column, player))

        if depth == 0 or legal_moves is None or self.game_over(state):
            return self.value(state, player)

        if player == self.max_player:
            for move in legal_moves:
                alpha = max(alpha, -self.search(move, depth-1, opp,
                            alpha, beta))
                if beta <= alpha:
                    break
            return alpha
        else:
            for move in legal_moves:
                beta = min(beta, -self.search(move, depth-1, opp, alpha, beta))
                if beta <= alpha:
                    break
            return beta

    def check_for_streak(self, state, color, streak):
        streaks = 0
        for col in range(self.board.width):
            for row in range(self.board.height):
                if state[col][row] == color:
                    streaks += self.vertical_streak(col, row, state, streak)
                    streaks += self.horizontal_streak(col, row, state, streak)
                    streaks += self.diagonal_streak(col, row, state, streak)
        return streaks

    def game_over(self, state):
        return (self.check_for_streak(state, self.colors[0], 4) >= 1 or
                self.check_for_streak(state, self.colors[1], 4) >= 1)

    def value(self, state, color):
        player = color
        opponent = self._other_color(player)
        my_fours = self.check_for_streak(state, player, 4)
        my_threes = self.check_for_streak(state, player, 3)
        my_twos = self.check_for_streak(state, player, 2)
        opp_fours = self.check_for_streak(state, opponent, 4)
        opp_threes = self.check_for_streak(state, opponent, 3)
        opp_twos = self.check_for_streak(state, opponent, 2)
        if opp_fours > 0:
            return -sys.maxsize
        else:
            return (my_fours*10000000 + my_threes*1000 + my_twos*10
                    - opp_threes*10000 - opp_twos*100)

    def vertical_streak(self, icol, irow, state, streak):
        count = 0
        for row in range(irow, self.board.height):
            if state[icol][row] == state[icol][irow]:
                count += 1
            else:
                break
        for row in range(irow, -1, -1):
            if state[icol][row] == state[icol][irow]:
                count += 1
            else:
                break
        return int(count >= streak)

    def horizontal_streak(self, icol, irow, state, streak):
        count = 0
        for col in range(icol, self.board.width):
            if state[col][irow] == state[icol][irow]:
                count += 1
            else:
                break
        for col in range(icol, -1, -1):
            if state[col][irow] == state[icol][irow]:
                count += 1
            else:
                break
        return int(count >= streak)

    def diagonal_streak(self, icol, irow, state, streak):
        total = 0
        # check for diagonals with positive slope
        count = 0
        col = icol
        for row in range(irow, self.board.height):
            if col >= self.board.width or col < 0:
                break
            if state[col][row] == state[icol][irow]:
                count += 1
            else:
                break
            col += 1  # increment column when row is incremented
        total += int(count >= streak)

        # check for diagonals with negative slope
        count = 0
        col = icol
        for row in range(irow, -1, -1):
            if col >= self.board.width or col < 0:
                break
            elif state[col][row] == state[icol][irow]:
                count += 1
            else:
                break
            col += 1  # increment column when row is incremented
        total += int(count >= streak)

        return total
<commit_msg>Finalized evaluation heuristic. Auto now results in draw.<commit_after>import operator
import sys

import error

class Minimax():
    def __init__(self, board, colors):
        self.board = board
        self.colors = colors
        self.max_player = None

    def _other_color(self, color):
        colors = list(self.colors)
        colors.remove(color)
        return colors[0]

    def best_move(self, color, difficulty):
        player = color
        self.max_player = player
        opponent = self._other_color(color)
        depth = difficulty

        legal_moves = {}
        for column in range(self.board.width):
            if (self.board.valid_move(column, self.board.get_board())
                is not None):
                state = self.make_fake_move(self.board.get_board(),
                                            column, player)
                legal_moves[column] = -self.search(state, depth-1, opponent)
        return max(legal_moves.items(), key=operator.itemgetter(1))

    def make_fake_move(self, state, column, color):
        row = self.board.valid_move(column, state)
        if row is None:
            raise error.InvalidMoveError
        state[column][row] = color
        return state

    def search(self, state, depth, player,
               alpha=-sys.maxsize, beta=sys.maxsize):
        opp = self._other_color(player)
        legal_moves = []
        for column in range(self.board.width):
            if self.board.valid_move(column, state) is not None:
                legal_moves.append(self.make_fake_move(state, column, player))

        if depth == 0 or legal_moves is None or self.game_over(state):
            return self.value(state, player)

        if player == self.max_player:
            for move in legal_moves:
                alpha = max(alpha, -self.search(move, depth-1, opp,
                            alpha, beta))
                if beta <= alpha:
                    break
            return alpha
        else:
            for move in legal_moves:
                beta = min(beta, -self.search(move, depth-1, opp, alpha, beta))
                if beta <= alpha:
                    break
            return beta

    def check_for_streak(self, state, color, streak):
        streaks = 0
        for col in range(self.board.width):
            for row in range(self.board.height):
                if state[col][row] == color:
                    streaks += self.vertical_streak(col, row, state, streak)
                    streaks += self.horizontal_streak(col, row, state, streak)
                    streaks += self.diagonal_streak(col, row, state, streak)
        return streaks

    def game_over(self, state):
        return (self.check_for_streak(state, self.colors[0], 4) >= 1 or
                self.check_for_streak(state, self.colors[1], 4) >= 1)

    def value(self, state, color):
        player = color
        opponent = self._other_color(player)
        fours = self.check_for_streak(state, player, 4)
        threes = self.check_for_streak(state, player, 3)
        twos = self.check_for_streak(state, player, 2)
        losses = self.check_for_streak(state, opponent, 4)
        if losses > 0:
            return -sys.maxsize
        else:
            return (fours*16 + threes*4 + twos*2)

    def vertical_streak(self, icol, irow, state, streak):
        count = 0
        for row in range(irow, self.board.height):
            if state[icol][row] == state[icol][irow]:
                count += 1
            else:
                break
        for row in range(irow, -1, -1):
            if state[icol][row] == state[icol][irow]:
                count += 1
            else:
                break
        return int(count >= streak)

    def horizontal_streak(self, icol, irow, state, streak):
        count = 0
        for col in range(icol, self.board.width):
            if state[col][irow] == state[icol][irow]:
                count += 1
            else:
                break
        for col in range(icol, -1, -1):
            if state[col][irow] == state[icol][irow]:
                count += 1
            else:
                break
        return int(count >= streak)

    def diagonal_streak(self, icol, irow, state, streak):
        total = 0
        # check for diagonals with positive slope
        count = 0
        col = icol
        for row in range(irow, self.board.height):
            if col >= self.board.width or col < 0:
                break
            if state[col][row] == state[icol][irow]:
                count += 1
            else:
                break
            col += 1  # increment column when row is incremented
        total += int(count >= streak)

        # check for diagonals with negative slope
        count = 0
        col = icol
        for row in range(irow, -1, -1):
            if col >= self.board.width or col < 0:
                break
            elif state[col][row] == state[icol][irow]:
                count += 1
            else:
                break
            col += 1  # increment column when row is incremented
        total += int(count >= streak)

        return total
<|endoftext|>


<commit_before>"""
Main API for Vega-lite spec generation
"""

try:
    import traitlets as T
except ImportError:
    from IPython.utils import traitlets as T

from .utils import parse_shorthand, infer_vegalite_type
from ._py3k_compat import string_types

import pandas as pd

class BaseObject(T.HasTraits):

    def __contains__(self, key):
        value = getattr(self, key)
        if isinstance(value, pd.DataFrame):
            return True
        return (value is not None) and (not (not isinstance(value, bool) and not value))

    def to_dict(self):
        result = {}
        for k in self.traits():
            if k in self:
                v = getattr(self, k)
                if isinstance(v, BaseObject):
                    result[k] = v.to_dict()
                else:
                    result[k] = v
        return result

    def __repr__(self):
        return repr(self.to_dict())

class Data(BaseObject):

    formatType = T.Enum(['json','csv'], default_value='json')
    url = T.Unicode(default_value=None, allow_none=True)
    values = T.List(default_value=None, allow_none=True)

    def to_dict(self):
        result = {'formatType':self.formatType,
                  'values':self.data.to_dict('records')
                  }
        return result

class Scale(BaseObject):
    pass

class Axis(BaseObject):
    pass

class Band(BaseObject):
    pass

class Legend(BaseObject):
    pass

class SortItems(BaseObject):
    name = T.Unicode(default_value=None, allow_none=True)
    aggregate = T.Enum(['avg','sum','min','max','count'], default_value=True)
    reverse = T.Bool(False)

class Shelf(BaseObject):
    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(Shelf, self).__init__(self, **kwargs)

    def _infer_type(self, data):
        if self.type is None and self.name in data:
            self.type = infer_vegalite_type(data[self.name])

    def _shorthand_changed(self, name, old, new):
        D = parse_shorthand(self.shorthand)
        for key, val in D.items():
            setattr(self, key, val)

    shorthand = T.Unicode('')
    name = T.Unicode('')
    type = T.Enum(['N', 'O', 'Q', 'T'], default_value=None, allow_none=True)
    timeUnit = T.Enum(['year', 'month', 'day', 'date',
                       'hours', 'minutes', 'seconds'],
                      default_value=None, allow_none=True)
    bin = T.Union([T.Bool(), T.Int()], default_value=False)
    sort = T.List(T.Instance(SortItems), default_value=None, allow_none=True)
    aggregate = T.Enum(['avg', 'sum', 'median', 'min', 'max', 'count'],
                       default_value=None, allow_none=True)

class DataCoordinate(Shelf):
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    axis = T.Instance(Axis, default_value=None, allow_none=True)
    band = T.Instance(Band, default_value=None, allow_none=True)

class X(DataCoordinate):
    pass

class Y(DataCoordinate):
    pass

class FacetCoordinate(Shelf):
    aggregate = T.Enum(['count'], default_value=None, allow_none=True)
    padding = T.CFloat(0.1)
    axis = T.Instance(Axis, default_value=None, allow_none=True)
    height = T.CInt(150)

class Row(FacetCoordinate):
    pass

class Col(FacetCoordinate):
    pass

class Size(Shelf):
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    value = T.CInt(30)

class Color(Shelf):
    value = T.Unicode('#4682b4')
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    opacity = T.Float(1.0)

class Shape(Shelf):
    value = T.Enum(['circle', 'square', 'cross', 'diamond', 'triangle-up',
                    'triangle-down'], default_value='circle')
    aggregate = T.Enum(['count'], default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    filled = T.Bool(False)

class Encoding(BaseObject):
    x = T.Union([T.Instance(X), T.Unicode()],
                default_value=None, allow_none=True)
    y = T.Union([T.Instance(Y), T.Unicode()],
                default_value=None, allow_none=True)
    row = T.Union([T.Instance(Row), T.Unicode()],
                  default_value=None, allow_none=True)
    col = T.Union([T.Instance(Col), T.Unicode()],
                  default_value=None, allow_none=True)
    size = T.Union([T.Instance(Size), T.Unicode()],
                   default_value=None, allow_none=True)
    color = T.Union([T.Instance(Color), T.Unicode()],
                    default_value=None, allow_none=True)
    shape = T.Union([T.Instance(Shape), T.Unicode()],
                    default_value=None, allow_none=True)
    parent = T.Instance(BaseObject, default_value=None, allow_none=True)

    def _infer_types(self, data):
        for attr in ['x', 'y', 'row', 'col', 'size', 'color', 'shape']:
            val = getattr(self, attr)
            if val is not None:
                val._infer_type(data)

    def _x_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.x = X(new)
        if getattr(self.parent, 'data', None) is not None:
            self.x._infer_type(self.parent.data)

    def _y_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.y = Y(new)
        if getattr(self.parent, 'data', None) is not None:
            self.y._infer_type(self.parent.data)

    def _row_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.row = Row(new)
        if getattr(self.parent, 'data', None) is not None:
            self.row._infer_type(self.parent.data)

    def _col_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.col = Col(new)
        if getattr(self.parent, 'data', None) is not None:
            self.col._infer_type(self.parent.data)

    def _size_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.size = Size(new)
        if getattr(self.parent, 'data', None) is not None:
            self.size._infer_type(self.parent.data)

    def _color_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.color = Color(new)
        if getattr(self.parent, 'data', None) is not None:
            self.color._infer_type(self.parent.data)

    def _shape_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.shape = Shape(new)
        if self.parent is not None:
            self.shape._infer_type(self.parent.data)

class Viz(BaseObject):

    marktype = T.Enum(['point','tick','bar','line',
                     'area','circle','square','text'], default_value='point')
    encoding = T.Instance(Encoding, default_value=None, allow_none=True)

    _data = T.Instance(Data, default_value=None, allow_none=True)
    data = T.Any(default_value=None, allow_none=True)

    def _encoding_changed(self, name, old, new):
        if isinstance(new, Encoding):
            self.encoding.parent = self
            if self.data is not None:
                self.encoding._infer_types(self.data)

    def _data_changed(self, name, old, new):
        if not isinstance(new, pd.DataFrame):
            self.data = pd.DataFrame(new)
            return
        self._data = Data(data=new)
        if self.encoding is not None:
            self.encoding._infer_types(self.data)

    def __init__(self, data, **kwargs):
        kwargs['data'] = data
        super(Viz,self).__init__(self, **kwargs)

    def to_dict(self):
        result = {}
        for k in self.traits():
            if k in self and k!='data':
                v = getattr(self, k)
                if k=='_data':
                    k='data'
                if isinstance(v, BaseObject):
                    result[k] = v.to_dict()
                else:
                    result[k] = v
        return result

    def encode(self, **kwargs):
        self.encoding = Encoding(**kwargs)
        return self

    def mark(self, mt):
        self.marktype = mt
        return self

    def mark_point(self):
        return self.mark('point')

    def mark_tick(self):
        return self.mark('tick')

    def mark_bar(self):
        return self.mark('bar')

    def mark_line(self):
        return self.mark('line')

    def mark_area(self):
        return self.mark('area')

    def mark_circle(self):
        return self.mark('circle')

    def mark_square(self):
        return self.mark('square')

    def mark_text(self):
        return self.mark('text')
<commit_msg>Fixing recursion.<commit_after>"""
Main API for Vega-lite spec generation
"""

try:
    import traitlets as T
except ImportError:
    from IPython.utils import traitlets as T

from .utils import parse_shorthand, infer_vegalite_type
from ._py3k_compat import string_types

import pandas as pd

class BaseObject(T.HasTraits):

    def __contains__(self, key):
        value = getattr(self, key)
        if isinstance(value, pd.DataFrame):
            return True
        return (value is not None) and (not (not isinstance(value, bool) and not value))

    def to_dict(self):
        result = {}
        for k in self.traits():
            if k in self:
                v = getattr(self, k)
                if isinstance(v, BaseObject):
                    result[k] = v.to_dict()
                else:
                    result[k] = v
        return result

    def __repr__(self):
        return repr(self.to_dict())

class Data(BaseObject):

    formatType = T.Enum(['json','csv'], default_value='json')
    url = T.Unicode(default_value=None, allow_none=True)
    values = T.List(default_value=None, allow_none=True)

    def to_dict(self):
        result = {'formatType':self.formatType,
                  'values':self.data.to_dict('records')
                  }
        return result

class Scale(BaseObject):
    pass

class Axis(BaseObject):
    pass

class Band(BaseObject):
    pass

class Legend(BaseObject):
    pass

class SortItems(BaseObject):
    name = T.Unicode(default_value=None, allow_none=True)
    aggregate = T.Enum(['avg','sum','min','max','count'], default_value=True)
    reverse = T.Bool(False)

class Shelf(BaseObject):
    def __init__(self, shorthand, **kwargs):
        kwargs['shorthand'] = shorthand
        super(Shelf, self).__init__(self, **kwargs)

    def _infer_type(self, data):
        if self.type is None and self.name in data:
            self.type = infer_vegalite_type(data[self.name])

    def _shorthand_changed(self, name, old, new):
        D = parse_shorthand(self.shorthand)
        for key, val in D.items():
            setattr(self, key, val)

    shorthand = T.Unicode('')
    name = T.Unicode('')
    type = T.Enum(['N', 'O', 'Q', 'T'], default_value=None, allow_none=True)
    timeUnit = T.Enum(['year', 'month', 'day', 'date',
                       'hours', 'minutes', 'seconds'],
                      default_value=None, allow_none=True)
    bin = T.Union([T.Bool(), T.Int()], default_value=False)
    sort = T.List(T.Instance(SortItems), default_value=None, allow_none=True)
    aggregate = T.Enum(['avg', 'sum', 'median', 'min', 'max', 'count'],
                       default_value=None, allow_none=True)

class DataCoordinate(Shelf):
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    axis = T.Instance(Axis, default_value=None, allow_none=True)
    band = T.Instance(Band, default_value=None, allow_none=True)

class X(DataCoordinate):
    pass

class Y(DataCoordinate):
    pass

class FacetCoordinate(Shelf):
    aggregate = T.Enum(['count'], default_value=None, allow_none=True)
    padding = T.CFloat(0.1)
    axis = T.Instance(Axis, default_value=None, allow_none=True)
    height = T.CInt(150)

class Row(FacetCoordinate):
    pass

class Col(FacetCoordinate):
    pass

class Size(Shelf):
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    value = T.CInt(30)

class Color(Shelf):
    value = T.Unicode('#4682b4')
    scale = T.Instance(Scale, default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    opacity = T.Float(1.0)

class Shape(Shelf):
    value = T.Enum(['circle', 'square', 'cross', 'diamond', 'triangle-up',
                    'triangle-down'], default_value='circle')
    aggregate = T.Enum(['count'], default_value=None, allow_none=True)
    legend = T.Instance(Legend, default_value=None, allow_none=True)
    filled = T.Bool(False)

class Encoding(BaseObject):
    x = T.Union([T.Instance(X), T.Unicode()],
                default_value=None, allow_none=True)
    y = T.Union([T.Instance(Y), T.Unicode()],
                default_value=None, allow_none=True)
    row = T.Union([T.Instance(Row), T.Unicode()],
                  default_value=None, allow_none=True)
    col = T.Union([T.Instance(Col), T.Unicode()],
                  default_value=None, allow_none=True)
    size = T.Union([T.Instance(Size), T.Unicode()],
                   default_value=None, allow_none=True)
    color = T.Union([T.Instance(Color), T.Unicode()],
                    default_value=None, allow_none=True)
    shape = T.Union([T.Instance(Shape), T.Unicode()],
                    default_value=None, allow_none=True)
    parent = T.Instance(BaseObject, default_value=None, allow_none=True)

    def to_dict(self):
        result = {}
        for k in self.traits():
            if k in self and k not in ('parent'):
                v = getattr(self, k)
                if isinstance(v, BaseObject):
                    result[k] = v.to_dict()
                else:
                    result[k] = v
        return result

    def _infer_types(self, data):
        for attr in ['x', 'y', 'row', 'col', 'size', 'color', 'shape']:
            val = getattr(self, attr)
            if val is not None:
                val._infer_type(data)

    def _x_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.x = X(new)
        if getattr(self.parent, 'data', None) is not None:
            self.x._infer_type(self.parent.data)

    def _y_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.y = Y(new)
        if getattr(self.parent, 'data', None) is not None:
            self.y._infer_type(self.parent.data)

    def _row_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.row = Row(new)
        if getattr(self.parent, 'data', None) is not None:
            self.row._infer_type(self.parent.data)

    def _col_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.col = Col(new)
        if getattr(self.parent, 'data', None) is not None:
            self.col._infer_type(self.parent.data)

    def _size_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.size = Size(new)
        if getattr(self.parent, 'data', None) is not None:
            self.size._infer_type(self.parent.data)

    def _color_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.color = Color(new)
        if getattr(self.parent, 'data', None) is not None:
            self.color._infer_type(self.parent.data)

    def _shape_changed(self, name, old, new):
        if isinstance(new, string_types):
            self.shape = Shape(new)
        if self.parent is not None:
            self.shape._infer_type(self.parent.data)

class Viz(BaseObject):

    marktype = T.Enum(['point','tick','bar','line',
                     'area','circle','square','text'], default_value='point')
    encoding = T.Instance(Encoding, default_value=None, allow_none=True)

    _data = T.Instance(Data, default_value=None, allow_none=True)
    data = T.Any(default_value=None, allow_none=True)

    def _encoding_changed(self, name, old, new):
        if isinstance(new, Encoding):
            self.encoding.parent = self
            if self.data is not None:
                self.encoding._infer_types(self.data)

    def _data_changed(self, name, old, new):
        if not isinstance(new, pd.DataFrame):
            self.data = pd.DataFrame(new)
            return
        self._data = Data(data=new)
        if self.encoding is not None:
            self.encoding._infer_types(self.data)

    def __init__(self, data, **kwargs):
        kwargs['data'] = data
        super(Viz,self).__init__(self, **kwargs)

    def to_dict(self):
        result = {}
        for k in self.traits():
            if k in self and k!='data':
                v = getattr(self, k)
                if k=='_data':
                    k='data'
                if isinstance(v, BaseObject):
                    result[k] = v.to_dict()
                else:
                    result[k] = v
        return result

    def encode(self, **kwargs):
        self.encoding = Encoding(**kwargs)
        return self

    def mark(self, mt):
        self.marktype = mt
        return self

    def mark_point(self):
        return self.mark('point')

    def mark_tick(self):
        return self.mark('tick')

    def mark_bar(self):
        return self.mark('bar')

    def mark_line(self):
        return self.mark('line')

    def mark_area(self):
        return self.mark('area')

    def mark_circle(self):
        return self.mark('circle')

    def mark_square(self):
        return self.mark('square')

    def mark_text(self):
        return self.mark('text')
<|endoftext|>


<commit_before>from collections import defaultdict

from matplotlib import pyplot as plt
import matplotlib.markers as mmarkers
import matplotlib.cm as cm
import matplotlib.colors as mcolors
import pandas as pd
import numpy as np

from cycler import cycler

def _determine_col_name(agg_shelf, shelf):
    if 'bin' in agg_shelf and agg_shelf.bin:
        agg_coll = '{}_{}_bin'.format(agg_shelf.name, shelf)
    else:
        agg_coll = agg_shelf.name

    return agg_coll

def _vl_line(ax, encoding, data, rc):
    x_name = _determine_col_name(encoding.x, 'x')
    y_name = _determine_col_name(encoding.y, 'y')
    df = data

    if 'color' in encoding:
        data_itr = _do_color(encoding.color, data, rc)
    else:
        data_itr = [(data, {'color': 'k'})]

    rets = []
    for df, sty in data_itr:
        if 'aggregate' in encoding.x:
            _agg_cols = [c for c in rc.agg_cols if c != x_name]
            df = _do_aggregate(encoding.x, df, _agg_cols)

        if 'aggregate' in encoding.y:
            _agg_cols = [c for c in rc.agg_cols if c != y_name]
            df = _do_aggregate(encoding.y, df, _agg_cols)

        sty_dict = {}
        sty_dict.update(rc.pl_kw)
        sty_dict.update(sty)

        ln, = ax.plot(x_name, y_name, data=df, **sty_dict)
        rets.append(ln)
    return rets

def _vl_area(ax, encoding, data, rc):
    ln, = _vl_line(ax, encoding, data, rc)
    x_data = ln.get_xdata()
    y_data = ln.get_ydata()
    c = ln.get_color()
    area = ax.fill_between(x_data, y_data, color=c)
    return ln, area

def _vl_point(ax, encoding, data, rc):
    pl_kw = rc.pl_kw
    agg_cols = rc.agg_cols
    pl_kw.setdefault('linestyle', 'none')
    if 'shape' in encoding:
        data_itr = _do_shape(encoding.shape, data, rc)
    else:
        data_itr = [(data, {'marker': 'o', 'fillstyle': 'none'})]

    if 'color' in encoding:
        data_itr = _do_color_itr(encoding.color, data_itr, rc)

    lns = []
    for df, sty in data_itr:
        x_name = _determine_col_name(encoding.x, 'x')
        y_name = _determine_col_name(encoding.y, 'y')

        if 'aggregate' in encoding.x:
            _agg_cols = [c for c in agg_cols if c != x_name]
            df = _do_aggregate(encoding.x, df, _agg_cols)

        if 'aggregate' in encoding.y:
            _agg_cols = [c for c in agg_cols if c != y_name]
            df = _do_aggregate(encoding.y, df, _agg_cols)

        sty_dict = {}
        sty_dict.update(pl_kw)
        sty_dict.update(sty)

        ln = ax.plot(x_name, y_name,  data=df, **sty_dict)
        lns.extend(ln)

    return lns

def _vl_bar(ax, encoding, data, rc):
    return ax.bar(encoding.x.name, encoding.y.name, data=data, **rc.pl_kw)

def _do_shape(shape, data, rc):
    """Sort out how to do shape

    Given an encoding + a possibly reduced DataFrame, return
    an iterator of (gb_key, DataFrame), style kwarg
    """
    filled = shape.filled
    if filled:
        shape_dict = rc.filledshape
    else:
        shape_dict = rc.allshape
    shape_col = _determine_col_name(shape, 'shape')

    if shape.type == 'Q':
        dig_data, bin_edges = _digitize_col(shape, data)
        data['{}_shape_bin'.format(shape.name)] = dig_data

    if 'aggregate' in shape:
        _agg_cols = [c for c in rc.agg_cols if c != shape_col]
        data = _do_aggregate(shape, data, _agg_cols)

    fill_sty = {'fillstyle': 'full' if filled else 'none'}
    if filled:
        fill_sty['markeredgewidth'] = 0

    gb = data.groupby(shape_col)

    for df in gb:
        sty = shape_dict[df[0]]
        sty.update(fill_sty)

        yield df[1], sty

def _do_color_itr(color, data_itr, rc):

    for df, sty in data_itr:
        for ret in _do_color(color, df, rc):
            _df, _sty = ret
            out = {}
            out.update(sty)
            out.update(_sty)
            yield _df, out

def _do_color(color, data, rc):

    color_col = _determine_col_name(color, 'color')

    if 'aggregate' in color:
        _agg_cols = [c for c in rc.agg_cols if c != color_col]
        df = _do_aggregate(color, data, _agg_cols)

    gb = data.groupby(color_col)
    sty = {}
    for df in gb:
        c = rc.cmap(rc.norm(df[0]))
        sty['color'] = c
        yield df[1], sty

def _do_aggregate(shelf, data, by_keys):

    agg_method = shelf.aggregate

    binned = data.groupby(by_keys)
    agg_func_name = _AGG_MAP[agg_method]

    data = getattr(binned, agg_func_name)().reset_index()
    return data

def _do_binning(vls, data, bin_key, plot_kwargs):
    encoding = vls.encoding
    shelf = getattr(encoding, bin_key)
    dig, bin_edges = _digitize_col(shelf, data)
    data['{}_{}_bin'.format(shelf.name, bin_key)] = dig
    if vls.marktype == 'bar':
        plot_kwargs['width'] = np.mean(np.diff(bin_edges)) * .9
        plot_kwargs['align'] = 'center'

    return data

def _digitize_col(bin_encoding, data):
    x_name = bin_encoding.name
    d_min, d_max = data[x_name].min(), data[x_name].max()
    bin_count = bin_encoding.bin
    if isinstance(bin_count, bool):
        bin_count = 3
    bin_count = int(bin_count)

    # add one to bin count as we are generating edges here
    bin_edges = np.linspace(d_min, d_max, bin_count + 1, endpoint=True)
    centers = (bin_edges[1:] + bin_edges[:-1]) / 2
    dig = np.digitize(data[x_name], bin_edges, right=True) - 1
    valid_mask = (-1 < dig) * (dig < len(centers))
    ret = np.zeros_like(dig, dtype='float')
    ret[valid_mask] = centers[dig[valid_mask]]
    ret[~valid_mask] = np.nan

    return ret, bin_edges

class _RenderContext(object):
    def __init__(self, agg_cols):
        filled_shapes = mmarkers.MarkerStyle.filled_markers

        all_shapes = filled_shapes + ('x',  '4', '3', '+', '2', '1')

        cyl = cycler('marker', filled_shapes)
        finite_cy_iter = iter(cyl)
        self.filledshape = defaultdict(lambda: next(finite_cy_iter))

        cyl = cycler('marker', all_shapes)
        finite_cy_iter = iter(cyl)
        self.allshape = defaultdict(lambda: next(finite_cy_iter))
        self.pl_kw = {}
        self.agg_cols = agg_cols

        self.norm = None
        self.cmap = None

def render(vls, data=None):
    """Render a vega-lite spec using matplotlib

    Parameters
    ----------
    vls : dict
        A dictionary complying with the vega-lite spec

    Returns
    -------
    fig : Figure
        The figure created

    ax_list : list
        The list of axes created

    arts : dict
        Dictionary keyed on Axes of the artists created
    """
    plot_kwargs = {}
    encoding = vls.encoding
    if data is None:
        data = vls.data
    data = pd.DataFrame(data)
    shelves = ('row', 'col', 'shape', 'size', 'color', 'y', 'x')
    # TODO these seem to be missing from the api.py
    # 'detail', 'text')
    used_cols = list(set(getattr(encoding, k).name
                         for k in shelves if k in encoding))
    data = data[used_cols]

    x_binned = 'bin' in encoding.x and encoding.x.bin
    y_binned = 'bin' in encoding.y and encoding.y.bin
    if x_binned and y_binned:
        raise NotImplementedError("Double binning not done yet")
    agg_cols = []
    for sh in shelves:
        if sh not in encoding:
            continue
        shelf = getattr(encoding, sh)
        if 'bin' in shelf and shelf.bin:
            data = _do_binning(vls, data, sh, plot_kwargs)
        agg_cols.append(_determine_col_name(shelf, sh))
    data = data.dropna()

    agg_cols = list(set(agg_cols))

    rc = _RenderContext(agg_cols)
    rc.pl_kw.update(plot_kwargs)

    if 'color' in encoding:
        # TODO do their color loop up right
        rc.cmap = cm.get_cmap('viridis')
        # TODO get the scale from 'type'
        # TODO deal with raw-domain stuff
        if encoding.color.type == 'Q':
            _cname = _determine_col_name(encoding.color, 'color')
            cmin, cmax = data[_cname].min(), data[_cname].max()
            rc.norm = mcolors.Normalize(cmin, cmax)

    plot_func = _MARK_DISPATCHER[vls.marktype]

    has_row = 'row' in encoding
    has_col = 'col' in encoding

    ax_map = {}
    ax_list = None

    if has_col and has_row:
        # sort out the names with respect to binning
        col_name = _determine_col_name(encoding.col, 'col')
        row_name = _determine_col_name(encoding.row, 'row')
        if 'aggregate' in encoding.row:
            _agg_cols = [c for c in agg_cols if c != row_name]
            data = _do_aggregate(encoding.row, data, _agg_cols)

        if 'aggregate' in encoding.col:
            _agg_cols = [c for c in agg_cols if c != col_name]
            data = _do_aggregate(encoding.col, data,  _agg_cols)

        row_labels = data[row_name].unique()
        col_labels = data[col_name].unique()
        grid_keys = [(_['row'], _['col'])
                     for _ in (cycler('row', row_labels) *
                               cycler('col', col_labels))]
        col_num, row_num = len(col_labels), len(row_labels)
        facet_iter = data.groupby([row_name, col_name])
    elif has_col:
        col_name = _determine_col_name(encoding.col, 'col')

        if 'aggregate' in encoding.col:
            _agg_cols = [c for c in agg_cols if c != col_name]
            data = _do_aggregate(encoding.col, data, _agg_cols)

        col_labels = data[col_name].unique()
        col_num, row_num = len(col_labels), 1
        grid_keys = list(col_labels)
        facet_iter = data.groupby(col_name)
    elif has_row:
        row_name = _determine_col_name(encoding.row, 'row')
        if 'aggregate' in encoding.row:
            _agg_cols = [c for c in agg_cols if c != row_name]
            data = _do_aggregate(encoding.row, data, _agg_cols)

        row_labels = data[row_name].unique()
        col_num, row_num = 1, len(row_labels)
        grid_keys = list(row_labels)
        facet_iter = data.groupby(row_name)
    else:
        grid_keys = [None, ]
        col_num, row_num = 1, 1

        def _inner():
            yield None, data
        facet_iter = _inner()

    fig, ax_list = plt.subplots(row_num, col_num,
                                sharex=True, sharey=True,
                                squeeze=False)
    for k, ax in zip(grid_keys, ax_list.ravel()):
        ax_map[k] = ax
        ax.set_prop_cycle(cycler('color', 'k'))
        if 'x' in encoding and ax.rowNum == row_num - 1:
            ax.set_xlabel(encoding.x.name)
        if 'y' in encoding and ax.colNum == 0:
            ax.set_ylabel(encoding.y.name)

    rets = {}
    for k, df in facet_iter:
        ax = ax_map[k]
        _r = plot_func(ax, encoding, df, rc)
        rets[k] = _r
        if k:
            ax.set_title(repr(k))

    return rets, ax_map

_MARK_DISPATCHER = {'area': _vl_area,      # fill below line
                    'bar': _vl_bar,        # bar
                    'circle': None,        # ??
                    'line': _vl_line,      # line
                    'point': _vl_point,    # scatter
                    'square': None,        # ??
                    'text': None,          # ??
                    'tick': None}          # ??

_AGG_MAP = {"avg": 'mean',
            "sum": 'sum',
            "min": 'min',
            "max": 'max',
            "count": 'count'}
<commit_msg>MNT: fix due to change in API<commit_after>from collections import defaultdict

from matplotlib import pyplot as plt
import matplotlib.markers as mmarkers
import matplotlib.cm as cm
import matplotlib.colors as mcolors
import pandas as pd
import numpy as np

from cycler import cycler

def _determine_col_name(agg_shelf, shelf):
    if 'bin' in agg_shelf and agg_shelf.bin:
        agg_coll = '{}_{}_bin'.format(agg_shelf.name, shelf)
    else:
        agg_coll = agg_shelf.name

    return agg_coll

def _vl_line(ax, encoding, data, rc):
    x_name = _determine_col_name(encoding.x, 'x')
    y_name = _determine_col_name(encoding.y, 'y')
    df = data

    if 'color' in encoding:
        data_itr = _do_color(encoding.color, data, rc)
    else:
        data_itr = [(data, {'color': 'k'})]

    rets = []
    for df, sty in data_itr:
        if 'aggregate' in encoding.x:
            _agg_cols = [c for c in rc.agg_cols if c != x_name]
            df = _do_aggregate(encoding.x, df, _agg_cols)

        if 'aggregate' in encoding.y:
            _agg_cols = [c for c in rc.agg_cols if c != y_name]
            df = _do_aggregate(encoding.y, df, _agg_cols)

        sty_dict = {}
        sty_dict.update(rc.pl_kw)
        sty_dict.update(sty)

        ln, = ax.plot(x_name, y_name, data=df, **sty_dict)
        rets.append(ln)
    return rets

def _vl_area(ax, encoding, data, rc):
    ln, = _vl_line(ax, encoding, data, rc)
    x_data = ln.get_xdata()
    y_data = ln.get_ydata()
    c = ln.get_color()
    area = ax.fill_between(x_data, y_data, color=c)
    return ln, area

def _vl_point(ax, encoding, data, rc):
    pl_kw = rc.pl_kw
    agg_cols = rc.agg_cols
    pl_kw.setdefault('linestyle', 'none')
    if 'shape' in encoding:
        data_itr = _do_shape(encoding.shape, data, rc)
    else:
        data_itr = [(data, {'marker': 'o', 'fillstyle': 'none'})]

    if 'color' in encoding:
        data_itr = _do_color_itr(encoding.color, data_itr, rc)

    lns = []
    for df, sty in data_itr:
        x_name = _determine_col_name(encoding.x, 'x')
        y_name = _determine_col_name(encoding.y, 'y')

        if 'aggregate' in encoding.x:
            _agg_cols = [c for c in agg_cols if c != x_name]
            df = _do_aggregate(encoding.x, df, _agg_cols)

        if 'aggregate' in encoding.y:
            _agg_cols = [c for c in agg_cols if c != y_name]
            df = _do_aggregate(encoding.y, df, _agg_cols)

        sty_dict = {}
        sty_dict.update(pl_kw)
        sty_dict.update(sty)

        ln = ax.plot(x_name, y_name,  data=df, **sty_dict)
        lns.extend(ln)

    return lns

def _vl_bar(ax, encoding, data, rc):
    return ax.bar(encoding.x.name, encoding.y.name, data=data, **rc.pl_kw)

def _do_shape(shape, data, rc):
    """Sort out how to do shape

    Given an encoding + a possibly reduced DataFrame, return
    an iterator of (gb_key, DataFrame), style kwarg
    """
    filled = shape.filled
    if filled:
        shape_dict = rc.filledshape
    else:
        shape_dict = rc.allshape
    shape_col = _determine_col_name(shape, 'shape')

    if shape.type == 'Q':
        dig_data, bin_edges = _digitize_col(shape, data)
        data['{}_shape_bin'.format(shape.name)] = dig_data

    if 'aggregate' in shape:
        _agg_cols = [c for c in rc.agg_cols if c != shape_col]
        data = _do_aggregate(shape, data, _agg_cols)

    fill_sty = {'fillstyle': 'full' if filled else 'none'}
    if filled:
        fill_sty['markeredgewidth'] = 0

    gb = data.groupby(shape_col)

    for df in gb:
        sty = shape_dict[df[0]]
        sty.update(fill_sty)

        yield df[1], sty

def _do_color_itr(color, data_itr, rc):

    for df, sty in data_itr:
        for ret in _do_color(color, df, rc):
            _df, _sty = ret
            out = {}
            out.update(sty)
            out.update(_sty)
            yield _df, out

def _do_color(color, data, rc):

    color_col = _determine_col_name(color, 'color')

    if 'aggregate' in color:
        _agg_cols = [c for c in rc.agg_cols if c != color_col]
        df = _do_aggregate(color, data, _agg_cols)

    gb = data.groupby(color_col)
    sty = {}
    for df in gb:
        c = rc.cmap(rc.norm(df[0]))
        sty['color'] = c
        yield df[1], sty

def _do_aggregate(shelf, data, by_keys):

    agg_method = shelf.aggregate
    binned = data.groupby(by_keys)
    agg_func_name = _AGG_MAP[agg_method]

    data = getattr(binned, agg_func_name)().reset_index()
    return data

def _do_binning(vls, data, bin_key, plot_kwargs):
    encoding = vls.encoding
    shelf = getattr(encoding, bin_key)
    dig, bin_edges = _digitize_col(shelf, data)
    data['{}_{}_bin'.format(shelf.name, bin_key)] = dig
    if vls.marktype == 'bar':
        plot_kwargs['width'] = np.mean(np.diff(bin_edges)) * .9
        plot_kwargs['align'] = 'center'

    return data

def _digitize_col(bin_encoding, data):
    x_name = bin_encoding.name
    d_min, d_max = data[x_name].min(), data[x_name].max()
    bin_count = bin_encoding.bin
    if isinstance(bin_count, bool):
        bin_count = 3
    bin_count = int(bin_count.maxbins)
    # add one to bin count as we are generating edges here
    bin_edges = np.linspace(d_min, d_max, bin_count + 1, endpoint=True)
    centers = (bin_edges[1:] + bin_edges[:-1]) / 2
    dig = np.digitize(data[x_name], bin_edges, right=True) - 1
    valid_mask = (-1 < dig) * (dig < len(centers))
    ret = np.zeros_like(dig, dtype='float')
    ret[valid_mask] = centers[dig[valid_mask]]
    ret[~valid_mask] = np.nan

    return ret, bin_edges

class _RenderContext(object):
    def __init__(self, agg_cols):
        filled_shapes = mmarkers.MarkerStyle.filled_markers

        all_shapes = filled_shapes + ('x',  '4', '3', '+', '2', '1')

        cyl = cycler('marker', filled_shapes)
        finite_cy_iter = iter(cyl)
        self.filledshape = defaultdict(lambda: next(finite_cy_iter))

        cyl = cycler('marker', all_shapes)
        finite_cy_iter = iter(cyl)
        self.allshape = defaultdict(lambda: next(finite_cy_iter))
        self.pl_kw = {}
        self.agg_cols = agg_cols

        self.norm = None
        self.cmap = None

def render(vls, data=None):
    """Render a vega-lite spec using matplotlib

    Parameters
    ----------
    vls : dict
        A dictionary complying with the vega-lite spec

    Returns
    -------
    fig : Figure
        The figure created

    ax_list : list
        The list of axes created

    arts : dict
        Dictionary keyed on Axes of the artists created
    """
    plot_kwargs = {}
    encoding = vls.encoding
    if data is None:
        data = vls.data
    data = pd.DataFrame(data)
    shelves = ('row', 'col', 'shape', 'size', 'color', 'y', 'x')
    # TODO these seem to be missing from the api.py
    # 'detail', 'text')
    used_cols = list(set(getattr(encoding, k).name
                         for k in shelves if k in encoding))
    data = data[used_cols]

    x_binned = 'bin' in encoding.x and encoding.x.bin
    y_binned = 'bin' in encoding.y and encoding.y.bin
    if x_binned and y_binned:
        raise NotImplementedError("Double binning not done yet")
    agg_cols = []
    for sh in shelves:
        if sh not in encoding:
            continue
        shelf = getattr(encoding, sh)
        if 'bin' in shelf and shelf.bin:
            data = _do_binning(vls, data, sh, plot_kwargs)
        agg_cols.append(_determine_col_name(shelf, sh))
    data = data.dropna()

    agg_cols = list(set(agg_cols))

    rc = _RenderContext(agg_cols)
    rc.pl_kw.update(plot_kwargs)

    if 'color' in encoding:
        # TODO do their color loop up right
        rc.cmap = cm.get_cmap('viridis')
        # TODO get the scale from 'type'
        # TODO deal with raw-domain stuff
        if encoding.color.type == 'Q':
            _cname = _determine_col_name(encoding.color, 'color')
            cmin, cmax = data[_cname].min(), data[_cname].max()
            rc.norm = mcolors.Normalize(cmin, cmax)

    plot_func = _MARK_DISPATCHER[vls.marktype]

    has_row = 'row' in encoding
    has_col = 'col' in encoding

    ax_map = {}
    ax_list = None

    if has_col and has_row:
        # sort out the names with respect to binning
        col_name = _determine_col_name(encoding.col, 'col')
        row_name = _determine_col_name(encoding.row, 'row')
        if 'aggregate' in encoding.row:
            _agg_cols = [c for c in agg_cols if c != row_name]
            data = _do_aggregate(encoding.row, data, _agg_cols)

        if 'aggregate' in encoding.col:
            _agg_cols = [c for c in agg_cols if c != col_name]
            data = _do_aggregate(encoding.col, data,  _agg_cols)

        row_labels = data[row_name].unique()
        col_labels = data[col_name].unique()
        grid_keys = [(_['row'], _['col'])
                     for _ in (cycler('row', row_labels) *
                               cycler('col', col_labels))]
        col_num, row_num = len(col_labels), len(row_labels)
        facet_iter = data.groupby([row_name, col_name])
    elif has_col:
        col_name = _determine_col_name(encoding.col, 'col')

        if 'aggregate' in encoding.col:
            _agg_cols = [c for c in agg_cols if c != col_name]
            data = _do_aggregate(encoding.col, data, _agg_cols)

        col_labels = data[col_name].unique()
        col_num, row_num = len(col_labels), 1
        grid_keys = list(col_labels)
        facet_iter = data.groupby(col_name)
    elif has_row:
        row_name = _determine_col_name(encoding.row, 'row')
        if 'aggregate' in encoding.row:
            _agg_cols = [c for c in agg_cols if c != row_name]
            data = _do_aggregate(encoding.row, data, _agg_cols)

        row_labels = data[row_name].unique()
        col_num, row_num = 1, len(row_labels)
        grid_keys = list(row_labels)
        facet_iter = data.groupby(row_name)
    else:
        grid_keys = [None, ]
        col_num, row_num = 1, 1

        def _inner():
            yield None, data
        facet_iter = _inner()

    fig, ax_list = plt.subplots(row_num, col_num,
                                sharex=True, sharey=True,
                                squeeze=False)
    for k, ax in zip(grid_keys, ax_list.ravel()):
        ax_map[k] = ax
        ax.set_prop_cycle(cycler('color', 'k'))
        if 'x' in encoding and ax.rowNum == row_num - 1:
            ax.set_xlabel(encoding.x.name)
        if 'y' in encoding and ax.colNum == 0:
            ax.set_ylabel(encoding.y.name)

    rets = {}
    for k, df in facet_iter:
        ax = ax_map[k]
        _r = plot_func(ax, encoding, df, rc)
        rets[k] = _r
        if k:
            ax.set_title(repr(k))

    return rets, ax_map

_MARK_DISPATCHER = {'area': _vl_area,      # fill below line
                    'bar': _vl_bar,        # bar
                    'circle': None,        # ??
                    'line': _vl_line,      # line
                    'point': _vl_point,    # scatter
                    'square': None,        # ??
                    'text': None,          # ??
                    'tick': None}          # ??

_AGG_MAP = {"avg": 'mean',
            "sum": 'sum',
            "min": 'min',
            "max": 'max',
            "count": 'count'}
<|endoftext|>


<commit_before>#!/usr/bin/env python
import rospy
import motionplanner
import pathplanner
import location
import math
import pose
import mapParser
from gc_msgs.msg import MotionMsg  # for sending commands to motors
from gc_msgs.msg import BumpMsg    # for listening to bump sensors
from gc_msgs.msg import PoseMsg    # for listening to when the kinect sees a block
from gc_msgs.msg import ObstacleAheadMsg    # for listening to when the kinect sees a wall
from gc_msgs.msg import ObstacleMsg
import time
import random

#
# This node is the state machine that controls robot behavior
#
# TODO: call mapUpdater.make obstacle when

class RobotBrain(object):

#####################
## state methods ####
#####################

    # params: none
    # returns: void
    # starts all publishers, subscribers; loads map and fills associated variables
    def __init__(self):

        # defining variables related to robot's state
        self.robotState = 'wandering'               # the state of the robot; can be wandering, traveling, searching, consuming
                                               #    digesting, dispensing
        self.searchCount = 0                        # counter used to ensure that don't spend too long searching
        self.numBlocksCollected = 0                 # number of blocks that the robot has collected so far
        self.mapList = []                           # list of obstacles
        self.blockLocations = []                    # list of locations of blocks
        self.waypointsList = []                      # list of waypoints to current destination
        self.inHamperFlag = 0                       # flag used to indicate whether a block on the conveyor belt
                                               #  has made it to the hamper
        self.wanderCount = 0                        # while wandering, count up and turn
        self.currentPose = pose.Pose(0.0,0.0,0.0);
        self.previousWaypointPose = pose.Pose(0.0, 0.0, 0.0); #used for travel

        # constants
        # TODO: pick values well
        self.NUM_BLOCKS_NEEDED = 9                  # number of blocks needed to complete wall
        self.MAX_SEARCH_COUNT = 500                 # arbitrary value; should be tested and set
        self.POSITION_THRESHOLD = .005              # acceptable error to reaching a position. In meters.
        self.END_LOCATION = location.Location(0.0, 0.0);
        self.MAX_WANDER_COUNT = 30;
        self.KINECT_ERROR = .01;                    # uncertainty of the kinect in m
        self.ROBOT_RADIUS = .6                      # radius of robot in m

        # making publishers/subscribers
        self.motionPub = rospy.Publisher("command/Motors", MotionMsg);
        self.bumpSub = rospy.Subscriber('bumpData', BumpMsg, self.handleBumpMsg);
        self.blockSeenSub = rospy.Subscriber('blockSeen', PoseMsg, self.handleBlockSeenMsg);
        self.obstacleAheadSub = rospy.Subscriber('obstacleAhead', ObstacleAheadMsg, self.handleObstacleAheadMsg);

        # loading and processing map
        [self.blockLocations, self.mapList] = mapParser.parseMap('map.txt', self.currentPose);

        # TODO for debugging in wander
        self.blockLocations = []

        # objects
        self.motionPlanner =motionplanner.MotionPlanner();
        self.pathPlanner = pathplanner.PathPlanner(self.mapList, self.ROBOT_RADIUS);
        #TODO: mapUpdater

        # ROSPY methods
        rospy.init_node('robotbrain')
        rospy.on_shutdown(self.onShutdown)

        return

    # params: none
    # returns: none
    # Robot needs more blocks, but does not know where to find them.
    # perform ant-like motion; move forward. If hit an obstacle, turn right or left randomly
    # increment wanderCount; when it maxes out, spin around
    def wander(self):

        # if wanderCount has counted up enough, spin 360 degrees
        if (self.wanderCount >= self.MAX_WANDER_COUNT):

            # stop forward motion
            self.motionPlanner.stopWheels()

            # wanderCount holds at the max value while turning
            # TODO: that method no longer exists?
            doneTurning = self.motionPlanner.rotate360()
            if(doneTurning):
                self.wanderCount = 0

        # travel forward, incrementing wanderCount
        # when an obstacle is seen, handleObstacleAhead will be called to cause turning
        else:

            self.motionPlanner.translate(.01)
            time.sleep(.5)
            self.wanderCount +=1
        return

    # params: none
    # returns: none
    # Robot needs to move from it's current location to a destination
    # global waypoints contains a list of points the robot can travel to in straight-line paths
    def travel(self):

        # if have reached the current waypoint, stop and move to the next waypoint. If at end of wayponts, enter seatch
        if (self.getDistance(currentPose, waypointsList[0]) < self.POSITION_THRESHOLD):
            self.motionPlanner.stopWheels();
            self.previousWaypointPose = self.currentPose;

            # if there are more waypoints, travel to next one
            if len(self.waypointsList > 1):
                self.waypointsList = self.waypointsList[1:];

            # if there are no more waypoints, are at expected block location. start searching.
            else:
                self.robotState = 'searching';
                self.searchCount = 0;

        # if haven't reached the current waypoint, travel towards it
        else:
            # takes currentPose, destinationLoc, angVel, vel, startPose
            self.motionPlanner.travelTowards(currentPose, waypointsList[0], .01, .1, previousWaypointPose);
        return

    # params: none
    # returns: none
    # Robot has reached a place where it believes a block is
    # turn in slow circle looking for block. If time-out, reject block location and set course for a new one
    # search gets called repeatedly, so can switch states in the middle
    def search(self):

        self.motionPlanner.rotate(.1);  #TODO: check that this is a good speed

        # if haven't been searching too long, keep searching
        if(self.searchCount < self.MAX_SEARCH_COUNT):
            time.sleep(.5);            # should be moving slow enough that nothing passes in and out of view in this time
                                       # sleeping keeps count from getting crazy-high
            self.searchCount += 1;

        # if have been searching too long, give up on this block. Might not actually exist
        else:

            # if know the locations of more blocks, set path to them and go get them
            if (len(self.blockLocations) > 1):
                self.blockLocations = blockLocations[1:];
                self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.blockLocations[0])
                self.robotState = 'travelling'

            # if don't know the location of more blocks, wander
            else:
                self.robotState = 'wandering'
                self.wanderCount = 0
        return

    # params: none
    # returns: none
    # Robot has a block currently in its vision.
    # Move until block hits the bump sensors at the mouth of the conveyor belt
    def consume(self):

        # blockLocation is updatd by kinect to be accurate
        self.motionPlanner.travelTowards(self.currentPose, self.blockLocations[0], .01, .01, self.previousWaypointPose);
        return

    # params: none
    # returns: none
    # Robot has a block on the conveyor belt
    # Move conveyor belts so block is added to wall
    # this is an atomic action
    def digest(self):
        # TODO: this is a simple implementation that doesn't reverse belts if something gets stuck

        # stop wheel motors and start ramp motors to move block up ramp
        self.motionPlanner.stopWheels();
        self.motionPlanner.startEatingBelts();
        self.motionPlanner.setHamperAngle(10); # so block can fall down sucessfully
                                      # TODO: get a real value for this

        # wait for block to go up the ramp
        while(self.inHamperFlag < 1):
            time.sleep(1);   # wait for one second

        time.sleep(1);       # in case block needs to finish falling
        self.inHamperFlag = 0     # resetting flag

        # once block is in hamper, use hamper belts to push block into forming wall
        self.motionPlanner.setHamperAngle(0); # fully close hamper so conveyor belts can work
        self.motionPlanner.startHamperBelt();
        time.sleep(5)                    # wait for blocks to be pushed into right shape

        # turn off motors
        self.motionPlanner.stopConveyorBelt();

        # if know the locations of more blocks, set path to them and go get them
        if (len(self.blockLocations) > 1):
            self.blockLocations = self.blockLocations[1:];
            self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.blockLocations[0])
            self.robotState = 'travelling'
        # if don't know the location of more blocks, wander
        else:
            self.robotState = 'wandering'
            self.wanderCount = 0
        return

    # params: none
    # returns: void
    # Robot has finished gathering blocks; commands robot to travel to endLocation and open hatch
    def dispense(self):

        self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.END_LOCATION);

        while len(self.waypointsList) > 0:
            # if have reached the current waypoint, stop and move to the next waypoint.
            print 'waypoints', self.waypointsList
            if (self.getDistance(self.currentPose, self.waypointsList[0]) < self.POSITION_THRESHOLD):
                self.motionPlanner.stopWheels();
                self.previousWaypointPose = self.currentPose;

                # if there are more waypoints, travel to next one
                if len(self.waypointsList > 1):
                    self.waypointsList = self.waypointsList[1:];

                # otherwise, at end location
                else:
                    break

            # if haven't reached waypoint, keep travelling towards it
            else:
                self.motionPlanner.travelTowards(self.currentPose, self.waypointsList[0], .01, .1, self.previousWaypointPose);

            # have reached end location. Open hamper and drive away
            self.motionPlanner.setHamperAngle(math.pi/2.0);
            self.motionPlanner.translate(.01)
            time.sleep(.5)
            self.motionPlanner.stopWheels()
            return

##################
# Helper methods #
##################

    # params: pose: a Pose
    #         loc: a Location
    # returns: distance between them (absolute value)
    def getDistance(self, pose1, loc2):
        return math.sqrt((pose1.getX()-loc2.getX())**2 + (pose1.getY()-loc2.getY())**2)

    # params: loc1: a Location
    #         loc2: a Location
    # returns: distance between them (absolute value)
    def getDistance(self, pose1, loc2):
        print pose1
        print loc2
        return math.sqrt((pose1.getX()-loc2.getX())**2 + (pose1.getY()-loc2.getY())**2)

###############################
# Interrupt-handling methods ##
###############################

    # params: none
    # returns: none
    # based on which bump sensor went off, call appropriate respose
    def handleBumpMsg(self, msg):
        # declaring global variables

        # bump sensors 0 and 1 are at the mouth
        if (msg.bumpNumber <= 1):

            #if are currently searching and something hits these bump sensors, are eating the block
            if self.robotState == 'consuming':
                self.searchCount = 0;
                self.robotState = 'digesting';

        # bump sensor 2 is at the end of the conveyor belt. Indicates that partially done digesting
        elif (msg.bumpNumber == 2):
            if (self.robotState == 'digesting'):
                self.inHamperFlag = 1; #used by digesting method

            else:
                #msg = MotionMsg(); # defaults to 0
                self.motionPlanner.stopWheels();
                # not sure where other bump sensors on chassis will be;
                # should stop and back away if they are hit
                # TODO let bump sensors handle turning in wander state when see obstacle
        return

    # param: msg PoseMsg
    # return: none
    # takes in message from kinect indicating that a block has been seen
    def handleBlockSeenMsg(self, msg):

        newBlockLocation = location.Location(msg.xLoc, msg.yLoc)

        if self.robotState == 'wandering':
            self.motionPlanner.stopWheels();
            self.blockLocations.append(newBlockLocation)
            self.robotState = 'searching'
            self.searchCount = 0

        elif robotState == 'travelling':

            self.motionPlanner.stopWheels();

            # if have reache the block travelling towards, just enter search
            if (self.getDistance(newBlockLocation, self.blockLocations[0]) < self.KINECT_ERROR):
                self.robotState = 'searching'
                self.searchCount = 0

            # if found a new block while travelling, add it to blockLocations
            else:
                newblockLocations = [newBlockLocation]
                for i in range(len(self.blockLocations)):
                    newblockLocations.append(self.blockLocations[i])
                self.blockLocations = newblockLocations
                self.robotState = 'searching'

        # if consuming and see a new block
        elif (self.robotState == 'consuming') and (self.getDistance(newBlockLocation, self.blockLocations[0]) < self.KINECT_ERROR):
            # will continue consuming, but make this block the next block to get

            newBlockLocations = [self.blockLocations[0], newBlockLocation]
            for i in range(len(blockLocations)-1):
                newBlockLocations.append(self.blockLocations[i+1])
            self.blockLocations = newBlockLocations

        elif robotState == 'digesting':
            # continue digesting, bt make this next block to get

            newBlockLocations = [self.blockLocations[0], newBlockLocation]
            for i in range(len(blockLocations)-1):
                newBlockLocations.append(self.blockLocations[i+1])
            self.blockLocations = newBlockLocations

        return

    # param: msg
    # return: none
    # when an obstacle is directly ahead: stop. if travelling or depositing, wait and then remake path.
    # If wandering, turn randomly 90 deg right or left
    def handleObstacleAheadMsg(self, msg):

        # if robot is wandering, turn 90 deg left or right
        if self.robotState == 'wandering':
            self.motionPlanner.stopWheels()
            startAngle = self.currentPose.getAngle()

        # decide whether to turn right or left
            rand = random.random()
            if rand < 0.5:
                goalAngle = startAngle + math.pi/2.0
            else:
                goalAngle = startAngle - math.pi/2.0

                    # turn until turn is completed
            while (math.abs(self.currentPose.getAngle() - goalAngle) > self.motionPlanner.ANG_ERR):
                self.motionPlanner.rotateTowards(self.currentPose.getAngle(), goalAngle, .01, startAngle)

        # if not wandering, wait for mapUpdater to update with obstacle ahead and replan path
        else:
            time.sleep(2)
            self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.blockLocations[0])

        return

    # param: none
    # return: none
    # shutdown behavior for rospy
    def onShutdown(self):
        return

    # param: none
    # return: none
    # carried the main state machine
    def main(self):
        while(self.numBlocksCollected < self.NUM_BLOCKS_NEEDED):
            if (self.robotState == 'wandering'):
                self.wander()

            elif (self.robotState == 'traveling'):
                self.travel()

            elif (self.robotState == 'searching'):
                self.search()

            elif (self.robotState == 'consuming'):
                self.consume()

            elif (self.robotState == 'digesting'):
                self.digest()

            self.dispense()
            return

############
## Main ####
############

if __name__ == '__main__':

    try:

        robotbrain = RobotBrain()
#        robotbrain.main()

        while (True):
            #robotbrain.motionPlanner.translateTowards(robotbrain.currentPose, location.Location(1.0, 0), 0.10, pose.Pose(0.,0.,0.))
            robotbrain.motionPlanner.translate(.10);
        rospy.spin()          # keeps python from exiting until node is stopped
    except rospy.ROSInterruptException: pass
<commit_msg>debugging hal<commit_after>#!/usr/bin/env python
import rospy
import motionplanner
import pathplanner
import location
import math
import pose
import mapParser
from gc_msgs.msg import MotionMsg  # for sending commands to motors
from gc_msgs.msg import BumpMsg    # for listening to bump sensors
from gc_msgs.msg import PoseMsg    # for listening to when the kinect sees a block
from gc_msgs.msg import ObstacleAheadMsg    # for listening to when the kinect sees a wall
from gc_msgs.msg import ObstacleMsg
import time
import random

#
# This node is the state machine that controls robot behavior
#
# TODO: call mapUpdater.make obstacle when

class RobotBrain(object):

#####################
## state methods ####
#####################

    # params: none
    # returns: void
    # starts all publishers, subscribers; loads map and fills associated variables
    def __init__(self):

        # defining variables related to robot's state
        self.robotState = 'wandering'               # the state of the robot; can be wandering, traveling, searching, consuming
                                               #    digesting, dispensing
        self.searchCount = 0                        # counter used to ensure that don't spend too long searching
        self.numBlocksCollected = 0                 # number of blocks that the robot has collected so far
        self.mapList = []                           # list of obstacles
        self.blockLocations = []                    # list of locations of blocks
        self.waypointsList = []                      # list of waypoints to current destination
        self.inHamperFlag = 0                       # flag used to indicate whether a block on the conveyor belt
                                               #  has made it to the hamper
        self.wanderCount = 0                        # while wandering, count up and turn
        self.currentPose = pose.Pose(0.0,0.0,0.0);
        self.previousWaypointPose = pose.Pose(0.0, 0.0, 0.0); #used for travel

        # constants
        # TODO: pick values well
        self.NUM_BLOCKS_NEEDED = 9                  # number of blocks needed to complete wall
        self.MAX_SEARCH_COUNT = 500                 # arbitrary value; should be tested and set
        self.POSITION_THRESHOLD = .005              # acceptable error to reaching a position. In meters.
        self.END_LOCATION = location.Location(0.0, 0.0);
        self.MAX_WANDER_COUNT = 30;
        self.KINECT_ERROR = .01;                    # uncertainty of the kinect in m
        self.ROBOT_RADIUS = .6                      # radius of robot in m

        # making publishers/subscribers
        self.motionPub = rospy.Publisher("command/Motors", MotionMsg);
        self.bumpSub = rospy.Subscriber('bumpData', BumpMsg, self.handleBumpMsg);
        self.blockSeenSub = rospy.Subscriber('blockSeen', PoseMsg, self.handleBlockSeenMsg);
        self.obstacleAheadSub = rospy.Subscriber('obstacleAhead', ObstacleAheadMsg, self.handleObstacleAheadMsg);

        # loading and processing map
        [self.blockLocations, self.mapList] = mapParser.parseMap('map.txt', self.currentPose);

        # TODO for debugging in wander
        self.blockLocations = []

        # objects
        self.motionPlanner =motionplanner.MotionPlanner();
        self.pathPlanner = pathplanner.PathPlanner(self.mapList, self.ROBOT_RADIUS);
        #TODO: mapUpdater

        # ROSPY methods
        rospy.init_node('robotbrain')
        rospy.on_shutdown(self.onShutdown)

        # sending initial wheel 0 message to make wheel controller happy
        # (it freaks out if it gets a positive message after a long time step)
        self.motionPlanner.stopWheels()

        return

    # params: none
    # returns: none
    # Robot needs more blocks, but does not know where to find them.
    # perform ant-like motion; move forward. If hit an obstacle, turn right or left randomly
    # increment wanderCount; when it maxes out, spin around
    def wander(self):

        # if wanderCount has counted up enough, spin 360 degrees
        if (self.wanderCount >= self.MAX_WANDER_COUNT):

            # stop forward motion
            self.motionPlanner.stopWheels()

            # wanderCount holds at the max value while turning
            # TODO: that method no longer exists?
            doneTurning = self.motionPlanner.rotate360()
            if(doneTurning):
                self.wanderCount = 0

        # travel forward, incrementing wanderCount
        # when an obstacle is seen, handleObstacleAhead will be called to cause turning
        else:

            self.motionPlanner.translate(.01)
            time.sleep(.5)
            self.wanderCount +=1
        return

    # params: none
    # returns: none
    # Robot needs to move from it's current location to a destination
    # global waypoints contains a list of points the robot can travel to in straight-line paths
    def travel(self):

        # if have reached the current waypoint, stop and move to the next waypoint. If at end of wayponts, enter seatch
        if (self.getDistance(currentPose, waypointsList[0]) < self.POSITION_THRESHOLD):
            self.motionPlanner.stopWheels();
            self.previousWaypointPose = self.currentPose;

            # if there are more waypoints, travel to next one
            if len(self.waypointsList > 1):
                self.waypointsList = self.waypointsList[1:];

            # if there are no more waypoints, are at expected block location. start searching.
            else:
                self.robotState = 'searching';
                self.searchCount = 0;

        # if haven't reached the current waypoint, travel towards it
        else:
            # takes currentPose, destinationLoc, angVel, vel, startPose
            self.motionPlanner.travelTowards(currentPose, waypointsList[0], .01, .1, previousWaypointPose);
        return

    # params: none
    # returns: none
    # Robot has reached a place where it believes a block is
    # turn in slow circle looking for block. If time-out, reject block location and set course for a new one
    # search gets called repeatedly, so can switch states in the middle
    def search(self):

        self.motionPlanner.rotate(.1);  #TODO: check that this is a good speed

        # if haven't been searching too long, keep searching
        if(self.searchCount < self.MAX_SEARCH_COUNT):
            time.sleep(.5);            # should be moving slow enough that nothing passes in and out of view in this time
                                       # sleeping keeps count from getting crazy-high
            self.searchCount += 1;

        # if have been searching too long, give up on this block. Might not actually exist
        else:

            # if know the locations of more blocks, set path to them and go get them
            if (len(self.blockLocations) > 1):
                self.blockLocations = blockLocations[1:];
                self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.blockLocations[0])
                self.robotState = 'travelling'

            # if don't know the location of more blocks, wander
            else:
                self.robotState = 'wandering'
                self.wanderCount = 0
        return

    # params: none
    # returns: none
    # Robot has a block currently in its vision.
    # Move until block hits the bump sensors at the mouth of the conveyor belt
    def consume(self):

        # blockLocation is updatd by kinect to be accurate
        self.motionPlanner.travelTowards(self.currentPose, self.blockLocations[0], .01, .01, self.previousWaypointPose);
        return

    # params: none
    # returns: none
    # Robot has a block on the conveyor belt
    # Move conveyor belts so block is added to wall
    # this is an atomic action
    def digest(self):
        # TODO: this is a simple implementation that doesn't reverse belts if something gets stuck

        # stop wheel motors and start ramp motors to move block up ramp
        self.motionPlanner.stopWheels();
        self.motionPlanner.startEatingBelts();
        self.motionPlanner.setHamperAngle(10); # so block can fall down sucessfully
                                      # TODO: get a real value for this

        # wait for block to go up the ramp
        while(self.inHamperFlag < 1):
            time.sleep(1);   # wait for one second

        time.sleep(1);       # in case block needs to finish falling
        self.inHamperFlag = 0     # resetting flag

        # once block is in hamper, use hamper belts to push block into forming wall
        self.motionPlanner.setHamperAngle(0); # fully close hamper so conveyor belts can work
        self.motionPlanner.startHamperBelt();
        time.sleep(5)                    # wait for blocks to be pushed into right shape

        # turn off motors
        self.motionPlanner.stopConveyorBelt();

        # if know the locations of more blocks, set path to them and go get them
        if (len(self.blockLocations) > 1):
            self.blockLocations = self.blockLocations[1:];
            self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.blockLocations[0])
            self.robotState = 'travelling'
        # if don't know the location of more blocks, wander
        else:
            self.robotState = 'wandering'
            self.wanderCount = 0
        return

    # params: none
    # returns: void
    # Robot has finished gathering blocks; commands robot to travel to endLocation and open hatch
    def dispense(self):

        self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.END_LOCATION);

        while len(self.waypointsList) > 0:
            # if have reached the current waypoint, stop and move to the next waypoint.
            print 'waypoints', self.waypointsList
            if (self.getDistance(self.currentPose, self.waypointsList[0]) < self.POSITION_THRESHOLD):
                self.motionPlanner.stopWheels();
                self.previousWaypointPose = self.currentPose;

                # if there are more waypoints, travel to next one
                if len(self.waypointsList > 1):
                    self.waypointsList = self.waypointsList[1:];

                # otherwise, at end location
                else:
                    break

            # if haven't reached waypoint, keep travelling towards it
            else:
                self.motionPlanner.travelTowards(self.currentPose, self.waypointsList[0], .01, .1, self.previousWaypointPose);

            # have reached end location. Open hamper and drive away
            self.motionPlanner.setHamperAngle(math.pi/2.0);
            self.motionPlanner.translate(.01)
            time.sleep(.5)
            self.motionPlanner.stopWheels()
            return

##################
# Helper methods #
##################

    # params: pose: a Pose
    #         loc: a Location
    # returns: distance between them (absolute value)
    def getDistance(self, pose1, loc2):
        return math.sqrt((pose1.getX()-loc2.getX())**2 + (pose1.getY()-loc2.getY())**2)

    # params: loc1: a Location
    #         loc2: a Location
    # returns: distance between them (absolute value)
    def getDistance(self, pose1, loc2):
        print pose1
        print loc2
        return math.sqrt((pose1.getX()-loc2.getX())**2 + (pose1.getY()-loc2.getY())**2)

###############################
# Interrupt-handling methods ##
###############################

    # params: none
    # returns: none
    # based on which bump sensor went off, call appropriate respose
    def handleBumpMsg(self, msg):
        # declaring global variables

        # bump sensors 0 and 1 are at the mouth
        if (msg.bumpNumber <= 1):

            #if are currently searching and something hits these bump sensors, are eating the block
            if self.robotState == 'consuming':
                self.searchCount = 0;
                self.robotState = 'digesting';

        # bump sensor 2 is at the end of the conveyor belt. Indicates that partially done digesting
        elif (msg.bumpNumber == 2):
            if (self.robotState == 'digesting'):
                self.inHamperFlag = 1; #used by digesting method

            else:
                #msg = MotionMsg(); # defaults to 0
                self.motionPlanner.stopWheels();
                # not sure where other bump sensors on chassis will be;
                # should stop and back away if they are hit
                # TODO let bump sensors handle turning in wander state when see obstacle
        return

    # param: msg PoseMsg
    # return: none
    # takes in message from kinect indicating that a block has been seen
    def handleBlockSeenMsg(self, msg):

        newBlockLocation = location.Location(msg.xLoc, msg.yLoc)

        if self.robotState == 'wandering':
            self.motionPlanner.stopWheels();
            self.blockLocations.append(newBlockLocation)
            self.robotState = 'searching'
            self.searchCount = 0

        elif robotState == 'travelling':

            self.motionPlanner.stopWheels();

            # if have reache the block travelling towards, just enter search
            if (self.getDistance(newBlockLocation, self.blockLocations[0]) < self.KINECT_ERROR):
                self.robotState = 'searching'
                self.searchCount = 0

            # if found a new block while travelling, add it to blockLocations
            else:
                newblockLocations = [newBlockLocation]
                for i in range(len(self.blockLocations)):
                    newblockLocations.append(self.blockLocations[i])
                self.blockLocations = newblockLocations
                self.robotState = 'searching'

        # if consuming and see a new block
        elif (self.robotState == 'consuming') and (self.getDistance(newBlockLocation, self.blockLocations[0]) < self.KINECT_ERROR):
            # will continue consuming, but make this block the next block to get

            newBlockLocations = [self.blockLocations[0], newBlockLocation]
            for i in range(len(blockLocations)-1):
                newBlockLocations.append(self.blockLocations[i+1])
            self.blockLocations = newBlockLocations

        elif robotState == 'digesting':
            # continue digesting, bt make this next block to get

            newBlockLocations = [self.blockLocations[0], newBlockLocation]
            for i in range(len(blockLocations)-1):
                newBlockLocations.append(self.blockLocations[i+1])
            self.blockLocations = newBlockLocations

        return

    # param: msg
    # return: none
    # when an obstacle is directly ahead: stop. if travelling or depositing, wait and then remake path.
    # If wandering, turn randomly 90 deg right or left
    def handleObstacleAheadMsg(self, msg):

        # if robot is wandering, turn 90 deg left or right
        if self.robotState == 'wandering':
            self.motionPlanner.stopWheels()
            startAngle = self.currentPose.getAngle()

        # decide whether to turn right or left
            rand = random.random()
            if rand < 0.5:
                goalAngle = startAngle + math.pi/2.0
            else:
                goalAngle = startAngle - math.pi/2.0

                    # turn until turn is completed
            while (math.abs(self.currentPose.getAngle() - goalAngle) > self.motionPlanner.ANG_ERR):
                self.motionPlanner.rotateTowards(self.currentPose.getAngle(), goalAngle, .01, startAngle)

        # if not wandering, wait for mapUpdater to update with obstacle ahead and replan path
        else:
            time.sleep(2)
            self.waypointsList = self.pathPlanner.plotPath(self.currentPose, self.blockLocations[0])

        return

    # param: none
    # return: none
    # shutdown behavior for rospy
    def onShutdown(self):
        return

    # param: none
    # return: none
    # carried the main state machine
    def main(self):
        while(self.numBlocksCollected < self.NUM_BLOCKS_NEEDED):
            if (self.robotState == 'wandering'):
                self.wander()

            elif (self.robotState == 'traveling'):
                self.travel()

            elif (self.robotState == 'searching'):
                self.search()

            elif (self.robotState == 'consuming'):
                self.consume()

            elif (self.robotState == 'digesting'):
                self.digest()

            self.dispense()
            return

############
## Main ####
############

if __name__ == '__main__':

    try:

        robotbrain = RobotBrain()
#        robotbrain.main()

        while (True):
            #robotbrain.motionPlanner.translateTowards(robotbrain.currentPose, location.Location(1.0, 0), 0.10, pose.Pose(0.,0.,0.))
            robotbrain.motionPlanner.rotate(-.5);
        rospy.spin()          # keeps python from exiting until node is stopped
    except rospy.ROSInterruptException: pass
<|endoftext|>