baseball_elimination.py

'''Code file for baseball elimination lab created for Advanced Algorithms
Spring 2020 at Olin College. The code for this lab has been adapted from:
https://github.com/ananya77041/baseball-elimination/blob/master/src/BaseballElimination.java'''

import sys
import math
import picos as pic
import networkx as nx
import itertools
import cvxopt
import matplotlib.pyplot as plt

class Division:
    '''
    The Division class represents a baseball division. This includes all the
    teams that are a part of that division, their winning and losing history,
    and their remaining games for the season.

    filename: name of a file with an input matrix that has info on teams &
    their games
    '''

    def __init__(self, filename):
        self.teams = {}
        self.G = nx.DiGraph()
        self.readDivision(filename)

    def readDivision(self, filename):
        '''Reads the information from the given file and builds up a dictionary
        of the teams that are a part of this division.

        filename: name of text file representing tournament outcomes so far
        & remaining games for each team
        '''
        f = open(filename, "r")
        lines = [line.split() for line in f.readlines()]
        f.close()

        lines = lines[1:]
        for ID, teaminfo in enumerate(lines):
            team = Team(int(ID), teaminfo[0], int(teaminfo[1]), int(teaminfo[2]), int(teaminfo[3]), list(map(int, teaminfo[4:])))
            self.teams[ID] = team

    def get_team_IDs(self):
        '''Gets the list of IDs that are associated with each of the teams
        in this division.

        return: list of IDs that are associated with each of the teams in the
        division
        '''
        return self.teams.keys()

    def is_eliminated(self, teamID, solver):
        '''Uses the given solver (either Linear Programming or Network Flows)
        to determine if the team with the given ID is mathematically
        eliminated from winning the division (aka winning more games than any
        other team) this season.

        teamID: ID of team that we want to check if it is eliminated
        solver: string representing whether to use the network flows or linear
        programming solver
        return: True if eliminated, False otherwise
        '''
        flag1 = False
        team = self.teams[teamID]

        temp = dict(self.teams)
        del temp[teamID]

        for _, other_team in temp.items():
            if team.wins + team.remaining < other_team.wins:
                print("Easy out", " ", team.name)
                flag1 = True

        saturated_edges = self.create_network(teamID)
        if not flag1:
            if solver == "Network Flows":
                flag1 = self.network_flows(saturated_edges)
            elif solver == "Linear Programming":
                flag1 = self.linear_programming(saturated_edges)
        print(self.teams[teamID].name,flag1)
        return flag1

    def draw_graph(self, graph, layout):
        """Draws a nice representation of a networkx graph object.
        Source: https://notebooks.azure.com/coells/projects/100days/html/day%2049%20-%20ford-fulkerson.ipynb"""

        plt.figure(figsize=(12, 4))
        plt.axis('off')

        nx.draw_networkx_nodes(graph, layout, node_color='steelblue', node_size=600)
        nx.draw_networkx_edges(graph, layout, edge_color='gray')
        nx.draw_networkx_labels(graph, layout, font_color='white')

        for u, v, e in graph.edges(data=True):
            label = '{}/{}'.format(e['flow'], e['capacity'])
            color = 'green' if e['flow'] < e['capacity'] else 'red'
            x = layout[u][0] * .6 + layout[v][0] * .4
            y = layout[u][1] * .6 + layout[v][1] * .4
            t = plt.text(x, y, label, size=16, color=color,
                         horizontalalignment='center', verticalalignment='center')

        plt.show()

    def create_network(self, THEteamID):
        '''Builds up the network needed for solving the baseball elimination
        problem as a network flows problem & stores it in self.G. Returns a
        dictionary of saturated edges that maps team pairs to the amount of
        additional games they have against each other.

        teamID: ID of team that we want to check if it is eliminated
        return: dictionary of saturated edges that maps team pairs to
        the amount of additional games they have against each other
        '''

        saturated_edges = {}
        self.G.clear()

        self.G.add_node('S')

        CapValue = self.teams[THEteamID].wins+self.teams[THEteamID].remaining

        self.G.add_node('T')

        a = self.get_team_IDs()
        for teamID in a:
            if(teamID != THEteamID):
                team = self.teams[teamID]
                teamWins = team.wins
                self.G.add_node(team.ID)
                self.G.add_edge(team.ID, 'T', capacity=CapValue-teamWins, flow= 0 )


        matchUps = []

        for m in range(0, len(self.teams)):
            teamID1 = m
            for n in range(m, len(self.teams)):
                teamID2 = n
                if(teamID1 != teamID2 and teamID1 != THEteamID and teamID2 != THEteamID):
                    matchUps.append((teamID1, teamID2))

        for match in matchUps:
            team1, team2 = match
            gamesR = self.teams[team1].get_against(team2)
            saturated_edges[match] = gamesR
            self.G.add_node(match)
            self.G.add_edge('S', match, capacity=gamesR, flow= 0 )
            for matchedTeam in match:
                self.G.add_edge(match, matchedTeam, capacity=sys.maxsize, flow=0)

        return saturated_edges


    def network_flows(self, saturated_edges):
        '''Uses network flows to determine if the team with given team ID
        has been eliminated. You can feel free to use the built in networkx
        maximum flow function or the maximum flow function you implemented as
        part of the in class implementation activity.

        saturated_edges: dictionary of saturated edges that maps team pairs to
        the amount of additional games they have against each other
        return: True if team is eliminated, False otherwise
        '''

        flow_val, flow_dict = nx.maximum_flow(self.G, 'S', 'T')


        for key in saturated_edges:
            currFlow=flow_dict['S'][key]
            if saturated_edges[key] != currFlow:
                return True

        return False

    def linear_programming(self, saturated_edges):
        '''Uses linear programming to determine if the team with given team ID
        has been eliminated. We recommend using a picos solver to solve the
        linear programming problem once you have it set up.
        Do not use the flow_constraint method that Picos provides (it does all of the work for you)
        We want you to set up the constraint equations using picos (hint: add_constraint is the method you want)

        saturated_edges: dictionary of saturated edges that maps team pairs to
        the amount of additional games they have against each other
        returns True if team is eliminated, False otherwise
        '''

        maxflow=pic.Problem()

        # Add the flow variables.
        f={}
        for edge in self.G.edges():
            upperVal = self.G[edge[0]][edge[1]]['capacity']
            if(upperVal < sys.maxsize):
                f[edge]=maxflow.add_variable('f[{0}]'.format(edge),1, lower=0, upper=self.G[edge[0]][edge[1]]['capacity'])
            else:
                f[edge]=maxflow.add_variable('f[{0}]'.format(edge),1, lower=0)


        # Add another variable for the total flow.
        F=maxflow.add_variable('F',1)

        s = 'S'
        t = 'T'
        # Enforce flow conservation.
        maxflow.add_list_of_constraints(
          [pic.sum([f[p,i] for p in self.G.predecessors(i)],'p','pred(i)')
            == pic.sum([f[i,j] for j in self.G.successors(i)],'j','succ(i)')
            for i in self.G.nodes() if i not in (s,t)],
          'i','nodes-(s,t)')

        # Set source flow at s.
        maxflow.add_constraint(
          pic.sum([f[p,s] for p in self.G.predecessors(s)],'p','pred(s)') + F
          == pic.sum([f[s,j] for j in self.G.successors(s)],'j','succ(s)'))

        # Set the objective.
        maxflow.set_objective('max',F)
        # Solve the problem.
        maxflow.solve(verbose=0,solver='cvxopt')

        for x in f:
            a,b = x
            if(a == 'S'):
                if abs(self.G[a][b]['capacity']-f[x].value)>1e-4:
                    return True
        return False


    def checkTeam(self, team):
        '''Checks that the team actually exists in this division.
        '''
        if team.ID not in self.get_team_IDs():
            raise ValueError("Team does not exist in given input.")

    def __str__(self):
        '''Returns pretty string representation of a division object.
        '''
        temp = ''
        for key in self.teams:
            temp = temp + f'{key}: {str(self.teams[key])} \n'
        return temp

class Team:
    '''
    The Team class represents one team within a baseball division for use in
    solving the baseball elimination problem. This class includes information
    on how many games the team has won and lost so far this season as well as
    information on what games they have left for the season.

    ID: ID to keep track of the given team
    teamname: human readable name associated with the team
    wins: number of games they have won so far
    losses: number of games they have lost so far
    remaining: number of games they have left this season
    against: dictionary that can tell us how many games they have left against
    each of the other teams
    '''

    def __init__(self, ID, teamname, wins, losses, remaining, against):
        self.ID = ID
        self.name = teamname
        self.wins = wins
        self.losses = losses
        self.remaining = remaining
        self.against = against

    def get_against(self, other_team=None):
        '''Returns number of games this team has against this other team.
        Raises an error if these teams don't play each other.
        '''
        try:
            num_games = self.against[other_team]
        except:
            raise ValueError("Team does not exist in given input.")

        return num_games

    def __str__(self):
        '''Returns pretty string representation of a team object.
        '''
        return f'{self.name} \t {self.wins} wins \t {self.losses} losses \t {self.remaining} remaining'

if __name__ == '__main__':
    if len(sys.argv) > 1:
        filename = sys.argv[1]
        division = Division(filename)
        for (ID, team) in division.teams.items():
            print(f'{team.name}: Eliminated? {division.is_eliminated(team.ID, "Linear Programming")}')
    else:
        print("To run this code, please specify an input file name. Example: python baseball_elimination.py teams2.txt.")