Skip to content

Commit

Permalink
Move control elements to sidebar and download trajectories figs
Browse files Browse the repository at this point in the history 
- this offers 3 figures.
- later we will change this and just show one figure
  • Loading branch information
@chraibi
chraibi committed Feb 28, 2024
1 parent 551ec23 commit 502cb6d
Show file tree
Hide file tree
Showing 3 changed files with 190 additions and 51 deletions.
60 changes: 25 additions & 35 deletions anim.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,7 @@ def _get_line_color(disk_color: str) -> str:
return "black" if brightness > 127 else "white"


def _create_orientation_line(
row: pd.DataFrame, line_length: float = 0.2, color: str = "black"
) -> Shape:
def _create_orientation_line(row: pd.DataFrame, line_length: float = 0.2, color: str = "black") -> Shape:
"""Create orientation Shape object."""
end_x = row["x"] + line_length * 0
end_y = row["y"] + line_length * 0
Expand Down Expand Up @@ -113,9 +111,7 @@ def _get_colormap(frame_data: pd.DataFrame, max_speed: float) -> List[Scatter]:
return [scatter_trace]


def _get_shapes_for_frame(
frame_data: pd.DataFrame, min_speed: float, max_speed: float
) -> Tuple[Shape, Scatter, Shape]:
def _get_shapes_for_frame(frame_data: pd.DataFrame, min_speed: float, max_speed: float) -> Tuple[Shape, Scatter, Shape]:
"""Construct circles as Shapes for agents, Hover and Directions."""

def create_shape(row: pd.DataFrame) -> Shape:
Expand Down Expand Up @@ -210,9 +206,7 @@ def _create_fig(
fig = go.Figure(
data=geometry_traces + initial_scatter_trace + initial_hover_trace,
frames=frames,
layout=go.Layout(
shapes=initial_shapes + initial_arrows, title=title, title_x=0.5
),
layout=go.Layout(shapes=initial_shapes + initial_arrows, title=title, title_x=0.5),
)
fig.update_layout(
updatemenus=[_get_animation_controls()],
Expand Down Expand Up @@ -275,9 +269,7 @@ def _get_slider_controls(steps: List[Dict[str, Any]]) -> Dict[str, Any]:
}


def _get_processed_frame_data(
data_df: pd.DataFrame, frame_num: int, max_agents: int
) -> Tuple[pd.DataFrame, int]:
def _get_processed_frame_data(data_df: pd.DataFrame, frame_num: int, max_agents: int) -> Tuple[pd.DataFrame, int]:
"""Process frame data and ensure it matches the maximum agent count."""
frame_data = data_df[data_df["frame"] == frame_num]
agent_count = len(frame_data)
Expand All @@ -302,40 +294,44 @@ def animate(
frames = data_df0["frame"].unique()
fr0 = frames.min()
fr1 = frames.max()
col1, col2, col3, col4, col5 = st.columns((5))
page_size = col5.number_input(
col1, col2, col3 = st.sidebar.columns((3))
p1 = col1.empty()
p2 = col2.empty()
p3 = col3.empty()
page_size = st.sidebar.number_input(
"Number of frames",
value=500,
min_value=100,
max_value=1000,
help="How many frames to animae. (the larger the slower)",
)
every_nth_frame = col4.number_input(
"fps",
value=16,
min_value=8,
max_value=100,
step=16,
help="Every nth frame.",
)
every_nth_frame = int(every_nth_frame)
with col1:
st.text("Backward")
p1.text("Backward")
decrement = st.button(":arrow_backward:")
if decrement:
datafactory.decrement_frame_start(int(page_size))
with col2:
st.text("Forward")
p2.text("Forward")
increment = st.button(":arrow_forward:")
if increment:
datafactory.increment_frame_start(int(page_size))

with col3:
st.text("Reset")
p3.text("Reset")
reset = st.button(":leftwards_arrow_with_hook:")
if reset:
datafactory.reset_frame_start(fr0)

every_nth_frame = st.sidebar.number_input(
"fps",
value=16,
min_value=8,
max_value=100,
step=16,
help="Show every nth frame.",
)
every_nth_frame = int(every_nth_frame)

if st.session_state.start_frame < fr0:
st.session_state.start_frame = fr0

Expand All @@ -346,9 +342,7 @@ def animate(
# Calculate page_end
frame_end = st.session_state.start_frame + page_size
frame_start = st.session_state.start_frame
data_df = data_df0[
(data_df0["frame"] >= frame_start) & (data_df0["frame"] <= frame_end)
]
data_df = data_df0[(data_df0["frame"] >= frame_start) & (data_df0["frame"] <= frame_end)]

min_speed = data_df["speed"].min()
max_speed = data_df["speed"].max()
Expand All @@ -367,12 +361,8 @@ def animate(
) = _get_shapes_for_frame(initial_frame_data, min_speed, max_speed)
color_map_trace = _get_colormap(initial_frame_data, max_speed)
for frame_num in selected_frames:
frame_data, agent_count = _get_processed_frame_data(
data_df, frame_num, max_agents
)
shapes, hover_traces, arrows = _get_shapes_for_frame(
frame_data, min_speed, max_speed
)
frame_data, agent_count = _get_processed_frame_data(data_df, frame_num, max_agents)
shapes, hover_traces, arrows = _get_shapes_for_frame(frame_data, min_speed, max_speed)
# title = f"<b>{title_note + ' | ' if title_note else ''}N: {agent_count}</b>"
title = f"<b>{title_note + ' | ' if title_note else ''}Number of Agents: {initial_agent_count}. Frame: {frame_num}</b>"
frame_name = str(int(frame_num))
Expand Down
155 changes: 145 additions & 10 deletions plots.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from plotly.graph_objs import Figure, Scatter
from plotly.subplots import make_subplots

plt.rcParams["text.usetex"] = True
st_column: TypeAlias = st.delta_generator.DeltaGenerator


Expand All @@ -34,11 +35,12 @@ def plot_trajectories(
data = trajectory_data.data
num_agents = len(np.unique(data["id"]))
colors = {
1: "magenta", # Assuming 1 is for female
2: "green", # Assuming 2 is for male
3: "black", # non binary
1: "green",
2: "purple",
3: "red",
4: "blue",
}
dnames = {1: "North", 2: "South", 3: "East", 4: "West"}
x_exterior, y_exterior = walkable_area.polygon.exterior.xy
x_exterior = list(x_exterior)
y_exterior = list(y_exterior)
Expand All @@ -48,7 +50,6 @@ def plot_trajectories(
if uid is not None:
df = data[data["id"] == uid]
direction = directions.loc[directions["id"] == uid, "direction_number"].iloc[0]

color_choice = colors[direction]
fig.add_trace(
go.Scatter(
Expand All @@ -60,6 +61,29 @@ def plot_trajectories(
name=f"ID {uid}",
)
)
fig.add_trace(
go.Scatter(
x=[df["x"].iloc[0]],
y=[df["y"].iloc[0]],
marker={"color": "red", "symbol": "circle"},
mode="markers",
name=f"Start ID {uid}",
)
)

for other, df in data.groupby("id"):
if other != uid:
fig.add_trace(
go.Scatter(
x=df["x"][::framerate],
y=df["y"][::framerate],
line={"color": "gray", "width": 0.1},
marker={"color": "gray"},
mode="lines",
name=f"ID {uid}",
)
)

else:
for uid, df in data.groupby("id"):
direction = directions.loc[directions["id"] == uid, "direction_number"].iloc[0]
Expand Down Expand Up @@ -96,25 +120,86 @@ def plot_trajectories(
x=x_exterior,
y=y_exterior,
mode="lines",
line={"color": "red"},
line={"color": "black"},
name="geometry",
)
)
count_direction = ""
ymin = -6
ymax = 4
for direction in [1, 2, 3, 4]:
count = directions[directions["direction_number"] == direction].shape[0]
count_direction += "Direction: " + str(direction) + ": " + str(count) + ". "
# count_direction += f"<span style='color:{colors[direction]};'>Direction</span> " + str(direction) + ": " + str(count) + ". "
count_direction += f"<span style='color:{colors[direction]};'> {dnames[direction]} {count}</span>."

fig.update_layout(
title=f" Trajectories: {num_agents}. {count_direction}",
xaxis_title="X",
yaxis_title="Y",
xaxis={"scaleanchor": "y"}, # range=[xmin, xmax]),
yaxis={"scaleratio": 1}, # range=[ymin, ymax]),
xaxis={"scaleanchor": "y"},
yaxis={"scaleratio": 1},
showlegend=False,
)
return fig


# mpl
def plot_trajectories_figure_mpl(
trajectory_data: pedpy.TrajectoryData,
walkable_area: pedpy.WalkableArea,
with_colors: bool,
):
"""Plot trajectories and geometry mpl version.
framerate: sampling rate of the trajectories.
"""
fig, ax = plt.subplots()
data = trajectory_data.data
num_agents = len(np.unique(data["id"]))
colors = {
1: "green",
2: "purple",
3: "red",
4: "blue",
}
dnames = {1: "North", 2: "South", 3: "East", 4: "West"}
x_exterior, y_exterior = walkable_area.polygon.exterior.xy
x_exterior = list(x_exterior)
y_exterior = list(y_exterior)

directions = assign_direction_number(data)
for uid, df in data.groupby("id"):
direction = directions.loc[directions["id"] == uid, "direction_number"].iloc[0]
if with_colors:
color_choice = colors[direction]
else:
color_choice = "gray"
plt.plot(
df["x"],
df["y"],
color=color_choice,
lw=0.1,
alpha=0.6,
)
# geometry
plt.plot(
x_exterior,
y_exterior,
color="black",
)
title_text = f"Trajectories: {num_agents}."
for direction in [1, 2, 3, 4]:
count = directions[directions["direction_number"] == direction].shape[0]
title_text += f" {dnames[direction]} {count}."

ax.set_title(title_text)
ax.set_xlabel(r"$x$\; /\;m")
ax.set_ylabel(r"$y$\; /\;m")

ax.set_aspect("equal", "box")
return fig


def plot_time_series(density: pd.DataFrame, speed: pd.DataFrame, fps: int) -> go.Figure:
"""Plot density and speed time series side-byside."""
fig = make_subplots(
Expand Down Expand Up @@ -158,7 +243,7 @@ def plot_time_series(density: pd.DataFrame, speed: pd.DataFrame, fps: int) -> go
fig.update_layout(
showlegend=False,
)
fig.update_yaxes(
fig.update_xaxes(
range=[rmin, rmax],
title_text=r"$\rho\; /\; 1/m^2$",
title_font={"size": 20},
Expand Down Expand Up @@ -228,7 +313,7 @@ def plot_fundamental_diagram_all(density_dict: Dict[str, pd.DataFrame], speed_di
"square",
"diamond",
"cross",
"x-thin",
"triangle-down",
"triangle-up",
"hexagon",
"star",
Expand Down Expand Up @@ -299,6 +384,56 @@ def plot_x_y(x: pd.Series, y: pd.Series, title: str, xlabel: str, ylabel: str, c
return trace, fig


def plot_fundamental_diagram_all_mpl(density_dict: dict, speed_dict: dict):
"""Plot fundamental diagram of all files using Matplotlib."""
# Define colors and marker styles
colors_const = [
"blue",
"red",
"green",
"magenta",
"black",
"cyan",
"yellow",
"orange",
"purple",
]
marker_shapes = [
"o",
"s",
"D",
"x",
"^",
"v",
"h",
"*",
"p",
] # Matplotlib marker styles

fig, ax = plt.subplots()
for i, ((filename, density), (_, speed)) in enumerate(zip(density_dict.items(), speed_dict.items())):
if isinstance(speed, pd.Series):
y = speed
else:
y = speed["speed"] # Adjust this if 'speed' DataFrame structure is different

ax.plot(
density["density"],
y,
color=colors_const[i % len(colors_const)],
alpha=0.5,
linestyle="",
marker=marker_shapes[i % len(marker_shapes)],
label=filename,
)

ax.set_xlabel(r"$\rho / m^{-2}$", fontsize=20)
ax.set_ylabel(r"$v\; / \frac{m}{s}$", fontsize=20)
ax.legend(loc="best")

return fig


def assign_direction_number(agent_data: pd.DataFrame) -> pd.DataFrame:
"""
Assign a direction number to each agent based on their main direction of motion.
Expand Down
Loading

0 comments on commit 502cb6d

Please sign in to comment.