Bachelor-Thesis_Furkan_Karakocaoglu/testumgebung/CrowdModelling/Assets/Depictions_Years/Scripts/Python/Heatmap_Multiple.py

import UnityEngine as ue
import numpy as np
import seaborn as sns
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.colors import LinearSegmentedColormap 
from matplotlib.tri import Triangulation

WIDTH = int(70)
HEIGHT = int(35)
OBSTACLE_PATH1 = "Assets/Data_image/obstacle1.pkl"
OBSTACLE_PATH2 = "Assets/Data_image/obstacle2.pkl"
OBSTACLE_PATH3 = "Assets/Data_image/obstacle3.pkl"
POSITION_PATH1 = ue.Application.dataPath + '/Data_position/Walk1.csv'
POSITION_PATH2 = ue.Application.dataPath + '/Data_position/Walk2.csv'
POSITION_PATH3 = ue.Application.dataPath + '/Data_position/Walk3.csv'
HEATMAP_PATH = "Assets/Data_image/heatmapMultiple1.png"

# Generate only if obstacles change
def set_obstacles(maximum, path):
    """
    set_obstacles creates a DataFrame which marks the positions of the obstacles. 
    The obstacles are marked with the negative maximum and the file is placed at the path.
    The negative maximum is used so that the cmap can assign the correct colors to the values.
    Since no other negative numbers exist and the next highest number is 0, the negative values are painted 'grey' 
    and the other [0-maximum] are painted according to the color gradient between the two specified colors.
    
    :param maximum: Is the maximum to the associated positions file
    :param path: Is the location of the file created in this function
    """
    global list_obstacle_length
    positions = pd.DataFrame(np.zeros((HEIGHT, WIDTH)))
    obstacles = ue.Object.FindObjectsOfType(ue.GameObject)
    for obstacle in obstacles:
        if(obstacle.layer == 15 or obstacle.layer == 12):
            startWidth = int(obstacle.transform.position.x - obstacle.transform.localScale.x / 2)
            endWidth = int(obstacle.transform.position.x + obstacle.transform.localScale.x / 2)
            startHeight = int(obstacle.transform.position.z - obstacle.transform.localScale.z / 2)
            endHeight = int(obstacle.transform.position.z + obstacle.transform.localScale.z / 2)
            for currentW in range(startWidth, endWidth, 1):
                for currentH in range(startHeight, endHeight, 1):
                    positions[currentW][currentH] = -maximum
    positions.to_pickle(path)

def set_patches(plt):
    """
    set_patches uses the positions of the obstacles to draw a black border around them. 
    
    :param plt: The plot, where the figure is placed, so it can be modified by this function
    """
    obstacles = ue.Object.FindObjectsOfType(ue.GameObject)
    for obstacle in obstacles:
        if(obstacle.layer == 15):
            startWidth = int(obstacle.transform.position.x - obstacle.transform.localScale.x / 2)
            endWidth = int(obstacle.transform.position.x + obstacle.transform.localScale.x / 2)
            startHeight = int(obstacle.transform.position.z - obstacle.transform.localScale.z / 2)
            endHeight = int(obstacle.transform.position.z + obstacle.transform.localScale.z / 2)
            # plt.gca().add_patch(
            plt.add_patch(
                patches.Rectangle(
                    (startWidth, startHeight),
                    endWidth - startWidth,
                    endHeight - startHeight,
                    fill=False,
                    color='black'
                )
            )

def triangulation_for_triheatmap(M, N):
    """
    triangulation_for_triheatmap divides each pixel of the heatmap into 4 triangles of equal size. 
    For this it first determines all x and y coordinates of the vertices of all triangles. 
    Then it creates the triangles based on the x and y coordinates.
    
    :param M: The width of the Heatmap
    :param N: The height of the Heatmap
    :return: The created triangles
    """
    xv, yv = np.meshgrid(np.arange(0, M+1), np.arange(0, N+1))  # vertices of the little squares
    xc, yc = np.meshgrid(np.arange(0.5, M), np.arange(0.5, N))  # centers of the little squares
    x = np.concatenate([xv.ravel(), xc.ravel()])
    y = np.concatenate([yv.ravel(), yc.ravel()])
    cstart = (M + 1) * (N + 1)  # indices of the centers

    trianglesN = [(i + 1 + (j + 1) * (M + 1), i + (j + 1) * (M + 1), cstart + i + j * M)
                  for j in range(N) for i in range(M)]
    trianglesE = [(i + 1 + j * (M + 1), i + 1 + (j + 1) * (M + 1), cstart + i + j * M)
                  for j in range(N) for i in range(M)]
    trianglesS = [(i + j * (M + 1), i + 1 + j * (M + 1), cstart + i + j * M)
                  for j in range(N) for i in range(M)]
    trianglesW = [(i + (j + 1) * (M + 1), i + j * (M + 1), cstart + i + j * M)
                  for j in range(N) for i in range(M)]
    
    return [Triangulation(x, y, triangles) for triangles in [trianglesN, trianglesE, trianglesS, trianglesW]]

# 1. Get position data from csv file
data1 = pd.read_csv(POSITION_PATH1, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
data1 = data1.round(0)
data2 = pd.read_csv(POSITION_PATH2, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
data2 = data2.round(0)
data3 = pd.read_csv(POSITION_PATH3, sep=';', usecols=["Position x", "Position z"], decimal=',', dtype={'Position x': float, 'Position z': float})
data3 = data3.round(0)

# 2. Group by positions and count appearance
data_count1 = data1.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
data_count2 = data2.groupby(['Position x', 'Position z']).size().reset_index(name='counts')
data_count3 = data3.groupby(['Position x', 'Position z']).size().reset_index(name='counts')

# 3. Create wide-form DataFrame for generating heatmap
positions1 = data_count1.loc[:,:].reset_index().pivot(index='Position z', columns='Position x', values='counts')
positions2 = data_count2.loc[:,:].reset_index().pivot(index='Position z', columns='Position x', values='counts')
positions3 = data_count3.loc[:,:].reset_index().pivot(index='Position z', columns='Position x', values='counts')

# 4. Fill missing values
positions1.fillna(0, inplace=True)
positions2.fillna(0, inplace=True)
positions3.fillna(0, inplace=True)

# 5. Reindex DataFrame (70,35) size of Surface; (70, 35) first x width then z height
positions1 = positions1.reindex_axis(range(0, HEIGHT), axis=0, fill_value=0)
positions1 = positions1.reindex_axis(range(0, WIDTH), axis=1, fill_value=0)
positions2 = positions2.reindex_axis(range(0, HEIGHT), axis=0, fill_value=0)
positions2 = positions2.reindex_axis(range(0, WIDTH), axis=1, fill_value=0)
positions3 = positions3.reindex_axis(range(0, HEIGHT), axis=0, fill_value=0)
positions3 = positions3.reindex_axis(range(0, WIDTH), axis=1, fill_value=0)

# 6.1 Get maximum value of positions data
max1 = pd.DataFrame(positions1).max().max() # Gets the maximum of each column and than the maximum of the maximums
max2 = pd.DataFrame(positions2).max().max() # Gets the maximum of each column and than the maximum of the maximums
max3 = pd.DataFrame(positions3).max().max() # Gets the maximum of each column and than the maximum of the maximums

# 6.2 Get obstacles (obstacles, market stalls) and paste specific value in positions
# Therefore save pkl file in folder and read from it afterwards
set_obstacles(max1, OBSTACLE_PATH1)
set_obstacles(max2, OBSTACLE_PATH2)
set_obstacles(max3, OBSTACLE_PATH3)
positions_obstacles = pd.read_pickle(OBSTACLE_PATH1)
positions_heatmap1 = pd.read_pickle(OBSTACLE_PATH1)
positions_heatmap2 = pd.read_pickle(OBSTACLE_PATH2)
positions_heatmap3 = pd.read_pickle(OBSTACLE_PATH3)

# 7 Merge positions data with obstacles data
positions_heatmap1.where(positions_heatmap1 != 0, positions1, inplace=True)
positions_heatmap2.where(positions_heatmap2 != 0, positions2, inplace=True)
positions_heatmap3.where(positions_heatmap3 != 0, positions3, inplace=True)

# 7.1 Debug Output 
# positions_heatmap1 = pd.DataFrame(positions_heatmap1)
# positions_heatmap1.to_html('Assets/Data_image/positions_heatmap1.html')
# positions_heatmap2 = pd.DataFrame(positions_heatmap2)
# positions_heatmap2.to_html('Assets/Data_image/positions_heatmap2.html')
# positions_heatmap3 = pd.DataFrame(positions_heatmap3)
# positions_heatmap3.to_html('Assets/Data_image/positions_heatmap3.html')

# 8.1 Divide the pixel in four equal triangles and plot heatmap
values = [positions_heatmap1, positions_heatmap2, positions_obstacles, positions_heatmap3]
triangul = triangulation_for_triheatmap(WIDTH, HEIGHT)
cmaps = [LinearSegmentedColormap.from_list(name='2019', colors=['grey', (0.40,0.76,0.65), (0.11,0.62,0.47)]),
         LinearSegmentedColormap.from_list(name='2020', colors=['grey', (0.99,0.55,0.38), (0.85,0.37,0.01)]),
         LinearSegmentedColormap.from_list(name='2021', colors=['grey', (1, 1, 0.975)]),
         LinearSegmentedColormap.from_list(name='2021', colors=['grey', (0.55,0.63,0.80), (0.46,0.44,0.70)])]
fig, ax = plt.subplots()
imgs = [ax.tripcolor(t, np.ravel(val), cmap=cmap)
        for t, val, cmap in zip(triangul, values, cmaps)]

# 8.2 Mark the Market stalls
set_patches(ax)

# 9.1 Set visual settings
ax.tick_params(axis='both', which='both', bottom='off', labelbottom='off', left='off', labelleft='off')
ax.margins(x=0, y=0)
ax.set_aspect('equal', 'box')  # square cells
plt.tight_layout()

# 9.2 Show Heatmap
plt.show()

# 10. Save Heatmap
fig.savefig(HEATMAP_PATH, transparent=True)