# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python3 (tf-gpu)
#     language: python
#     name: tf-gpu
# ---

import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import pickle
import numpy as np

bus_data_normal = pickle.load(open('output/bus_data_normal_traffic.pkl', 'rb'))
bus_data_remove_traffic = pickle.load(open('./bus_data_removed_background_equal_speed.pkl', 'rb'))

bus_id = 'Route16_trip174568020_06:30:00'

# +
# %matplotlib notebook

ax = plt.axes(projection='3d')

df_normal = pd.DataFrame(bus_data_normal[bus_id])
df_normal['sim'] = 'normal'
df_1m = pd.DataFrame(bus_data_remove_traffic[bus_id])
df_1m['sim'] = '1-min'

print(bus_id)
normal_traj = np.array(df_normal['position'].tolist())
remove_traj = np.array(df_1m['position'].tolist())

ax.scatter3D(remove_traj[:, 0], remove_traj[:, 1], df_1m['step'].tolist(), s=0.3, c='b', label='1-min', alpha=0.5)
ax.scatter3D(normal_traj[:, 0], normal_traj[:, 1], df_normal['step'].tolist(), s=0.3, c='orange', label='normal', alpha=0.5)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('time step', rotation = -90)

plt.legend()
plt.show()

# -



StarcoderdataJupyterScriptsDedupFiltered