From 502cb6d2da97b56b62ff780610095fc60c80ae4d Mon Sep 17 00:00:00 2001 From: Mohcine Chraibi Date: Wed, 28 Feb 2024 08:28:32 +0100 Subject: [PATCH] Move control elements to sidebar and download trajectories figs - this offers 3 figures. - later we will change this and just show one figure --- anim.py | 60 +++++++++----------- plots.py | 155 ++++++++++++++++++++++++++++++++++++++++++++++++---- traj_tab.py | 26 +++++++-- 3 files changed, 190 insertions(+), 51 deletions(-) diff --git a/anim.py b/anim.py index 168e613..ffdd900 100644 --- a/anim.py +++ b/anim.py @@ -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 @@ -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: @@ -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()], @@ -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) @@ -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 @@ -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() @@ -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"{title_note + ' | ' if title_note else ''}N: {agent_count}" title = f"{title_note + ' | ' if title_note else ''}Number of Agents: {initial_agent_count}. Frame: {frame_num}" frame_name = str(int(frame_num)) diff --git a/plots.py b/plots.py index 19f9baf..90e5d74 100644 --- a/plots.py +++ b/plots.py @@ -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 @@ -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) @@ -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( @@ -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] @@ -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"Direction " + str(direction) + ": " + str(count) + ". " + count_direction += f" {dnames[direction]} {count}." + 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( @@ -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}, @@ -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", @@ -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. diff --git a/traj_tab.py b/traj_tab.py index 3d54e8b..079fbd8 100644 --- a/traj_tab.py +++ b/traj_tab.py @@ -62,30 +62,32 @@ def run_tab2(selected_file: str, msg: DeltaGenerator) -> None: st.dataframe(trajectory_data.data.loc[:, columns_to_display]) if do_plot_trajectories: - c1, c2 = st.columns((0.8, 0.25)) - sample_frame = c2.slider( + st.sidebar.write("**Plot configuration**") + sample_frame = st.sidebar.slider( "Every nth frame", 1, 1000, 1, 10, - help="plot every_nth_frame", + help="plot every_nth_frame.", ) - uid = c2.number_input( + uid = st.sidebar.number_input( "Insert id of pedestrian", value=None, min_value=int(min(ids)), max_value=int(max(ids)), placeholder=f"Type a number in [{int(min(ids))}, {int(max(ids))}]", format="%d", + help="Visualize a single pedestrian.", ) - show_direction = c2.number_input( + show_direction = st.sidebar.number_input( "Choose direction to show", value=None, min_value=1, max_value=4, placeholder="Type a number in [1, 4]", format="%d", + help="Visualize pedestrians moving in a direction. **1: North. 2: South. 3: East. 4: West.**", ) fig = plots.plot_trajectories( trajectory_data, @@ -94,7 +96,19 @@ def run_tab2(selected_file: str, msg: DeltaGenerator) -> None: uid, show_direction, ) - c1.plotly_chart(fig) + st.plotly_chart(fig) + # matplotlib figs + c1, c2 = st.columns(2) + fig2 = plots.plot_trajectories_figure_mpl(trajectory_data, walkable_area, with_colors=True) + c1.pyplot(fig2) + figname = "trajectories_" + selected_file.split("/")[-1].split(".txt")[0] + "_colors.pdf" + fig2.savefig(figname) + plots.download_file(figname, c1, label="color") + fig3 = plots.plot_trajectories_figure_mpl(trajectory_data, walkable_area, with_colors=False) + c2.pyplot(fig3) + figname = "trajectories_" + selected_file.split("/")[-1].split(".txt")[0] + "_gray.pdf" + fig3.savefig(figname) + plots.download_file(figname, c2, label="gray") end_time = time.time() elapsed_time = end_time - start_time