Heatmap Generator

# @author: Brian Keschinger
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import pandas as pd

# Start utility functions
'''-------------------------------------------------------
Function: ResetPlot
Description: Clears the figure being used in pyplot and sets up our dimensions, axes, image, and layout.
----------------------------------------------------------'''
def ResetPlot():
    global ax
    global img
    # Clear the current plot so we can reuse it
    plt.clf()

    # Turn off auto scaling and set the limits
    ax = plt.gca()
    ax.set_autoscale_on(False)

    # hide both axes
    ax.axes.get_xaxis().set_visible(False)
    ax.axes.get_yaxis().set_visible(False)

    # Set axis limits [xmin, xmax, ymin, ymax]
    plt.axis([1, 1920, 1, 1080])

    # flipud(): flips an array in the "up/down" direction so it's flipping the image
    # Then we say that the plot origin is lower which also flips the image
    # The result is a correctly shown image with a coordinate system origin of the lower left
    plt.imshow(np.flipud(img), origin='lower')

    # We're hiding the axes and their labels anyway so this will auto fit nicely to the figure area
    plt.tight_layout()


'''----------------------------------------------------------
Function: ConvertDataToAdjustedGrid
Description: Used to convert the image's pixel-based grid to a new grid size
@param[in] dataFrame: Data Frame that contains a type column that contains kills and deaths
@param[in] cellSize: int that's the new grid cell size used to convert the 1920x1080 grid
@return[out] adjustedGrid: a 2D Array consisting of summed integer values that will act as the converted grid
-------------------------------------------------------------'''
def ConvertDataToAdjustedGrid(dataFrame, cellSize):
    # Using the gridSize find out how many rows and columns the new array and grid will be
    xAxisSize = int(1920 / cellSize)
    yAxisSize = int(1080 / cellSize)

    # Create a fixed-size array that starts with all zeros
    adjustedGrid = np.zeros((yAxisSize, xAxisSize))

    # Loop through all the rows of the data
    # Convert the raw pixel coordinates to the adjusted grid
    # Add 1 for any kill and subtract 1 for deaths
    for index, row in dataFrame.iterrows():
        xCell = int(np.floor(row.x / cellSize))
        yCell = (yAxisSize-1) - int(np.floor(row.y / cellSize)) # 0,0 for a 2D array will be the top left so we want to flip the yAxis
        adjustedGrid[yCell][xCell] += 1 if (row.type == 'k') else -1

    return adjustedGrid

'''----------------------------------------------------------
Function: CreateAndExportHeatmap
Description: Used to create and save to disk a heatmap from a dataset
@param[in] fileName: string that will be used to save the image to disc and name the graph
@param[in] dataFrame: Data Frame that contains a type column that contains kills and deaths
@param[in] cellSize: int that's the new grid cell size used to convert the 1920x1080 grid
-------------------------------------------------------------'''
def CreateAndExportHeatmap(fileName, dataFrame, cellSize):
    global ax
    ResetPlot()
    newAdjustedGrid = ConvertDataToAdjustedGrid(dataFrame, cellSize)

    # Find the min and max value in the array
    minVal = np.amin(newAdjustedGrid)
    maxVal = np.amax(newAdjustedGrid)

    # Between the min and max find our new bounds
    upperBound = max(abs(minVal), abs(maxVal))

    # Using the upperBound we found, set the color bounds from -upperBound to +upperBound so that white sits at 0 properly
    sns_heatmap = sns.heatmap(newAdjustedGrid, cmap='seismic_r', center=0.0, linewidths=.5, vmin=-upperBound, vmax=upperBound)

    # Set the plot title from the fileName but replace dashes with spaces
    ax.set_title(fileName.replace('-',' '))

    # Save the image to disk (directory must be created)
    # We're using 500 dpi because the original figure size is small and 500 will push the actual graph to be larger than 1920x1080
    # We'll pull this into PhotoShop and overlay it onto our map with some nice opacities, etc.
    fig = sns_heatmap.get_figure()
    fig.savefig('Exported_Graphs/' + fileName +'.png', dpi=500)
# End utility functions

# Read the CSVs from disk
CTF_data = pd.read_csv("Playtest_Data/6-7_Playtest_CTF.csv")
BMB_data = pd.read_csv("Playtest_Data/6-7_Playtest_Bomb.csv")

# data[key] == something returns a boolean value
# by insterting that value into data.loc it returns all the rows where that is true
CTF_deathPositions = CTF_data.loc[CTF_data['type']=='d']
CTF_killPositions = CTF_data.loc[CTF_data['type']=='k']

# Import Map Background
img = plt.imread("Source_Images/Map_Overhead.png")

# Because everything is plotted on an image of 1920x1080, we'll need a common factor to be the new cell size
# 1, 2, 3, 4, 5, 6, 8, 10, 12, 15, 20, 24, 30, 40, 60, 120
newGridCellSize = 120

CreateAndExportHeatmap('CTF-All-Kills-and-Deaths-Heatmap', CTF_data, newGridCellSize)
CreateAndExportHeatmap('BMB-All-Kills-and-Deaths-Heatmap', BMB_data, newGridCellSize)