notebooks/fundamental_diagram.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the steady state frames, used in the paper describing the methods. Also uses roughly the same style in the plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "steady_states = {\n",
    "    \"uo-050-180-180\": [211, 800],\n",
    "    \"uo-060-180-180\": [243, 771],\n",
    "    \"uo-070-180-180\": [203, 1113],\n",
    "    \"uo-100-180-180\": [200, 790],\n",
    "    \"uo-145-180-180\": [300, 1097],\n",
    "    \"uo-180-180-070\": [500, 1399],\n",
    "    \"uo-180-180-095\": [400, 1350],\n",
    "    \"uo-180-180-120\": [300, 1099],\n",
    "    \"uo-180-180-180\": [400, 1284],\n",
    "}\n",
    "\n",
    "style_options = {\n",
    "    \"uo-050-180-180\": {\"color\": \"red\", \"marker\": \"+\"},\n",
    "    \"uo-060-180-180\": {\"color\": \"green\", \"marker\": \"x\"},\n",
    "    \"uo-070-180-180\": {\"color\": \"blue\", \"marker\": \"x\"},\n",
    "    \"uo-100-180-180\": {\"color\": \"white\", \"marker\": \"s\", \"edgecolors\": \"pink\"},\n",
    "    \"uo-145-180-180\": {\"color\": \"cyan\", \"marker\": \"s\"},\n",
    "    \"uo-180-180-070\": {\"color\": \"grey\", \"marker\": \"^\"},\n",
    "    \"uo-180-180-095\": {\"color\": \"white\", \"marker\": \"^\", \"edgecolors\": \"orange\"},\n",
    "    \"uo-180-180-120\": {\n",
    "        \"color\": \"black\",\n",
    "        \"marker\": \"o\",\n",
    "    },\n",
    "    \"uo-180-180-180\": {\"color\": \"white\", \"marker\": \"o\", \"edgecolors\": \"purple\"},\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load trajectories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "\n",
    "from pedpy import TrajectoryData, TrajectoryUnit, load_trajectory\n",
    "\n",
    "folder = pathlib.Path(\"demo-data/uni-directional\")\n",
    "trajectories = {}\n",
    "trajectories_in_steady_state = {}\n",
    "for file in folder.glob(\"uo*.txt\"):\n",
    "    trajectory = load_trajectory(\n",
    "        trajectory_file=file,\n",
    "        default_frame_rate=16.0,\n",
    "        default_unit=TrajectoryUnit.CENTIMETER,\n",
    "    )\n",
    "    trajectory_in_steady_state = TrajectoryData(\n",
    "        data=trajectory.data[\n",
    "            trajectory.data.frame.between(\n",
    "                steady_states[file.stem][0], steady_states[file.stem][1]\n",
    "            )\n",
    "        ],\n",
    "        frame_rate=trajectory.frame_rate,\n",
    "    )\n",
    "\n",
    "    trajectories[file.stem] = trajectory\n",
    "    trajectories_in_steady_state[file.stem] = trajectory_in_steady_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from shapely import Polygon\n",
    "\n",
    "from pedpy import WalkableArea\n",
    "\n",
    "walkable_area = WalkableArea(\n",
    "    Polygon(\n",
    "        [\n",
    "            (2.8, -6.5),\n",
    "            (2.8, -4),\n",
    "            (1.8, -4),\n",
    "            (1.8, 4),\n",
    "            (2.8, 4),\n",
    "            (2.8, 8),\n",
    "            (-1, 8),\n",
    "            (-1, 4),\n",
    "            (0, 4),\n",
    "            (0, -4),\n",
    "            (-1, -4),\n",
    "            (-1, -6.5),\n",
    "        ]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define measurement areas and lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy import MeasurementArea, MeasurementLine\n",
    "\n",
    "measurement_area = MeasurementArea([(0, -2), (0, 0), (1.8, 0), (1.8, -2)])\n",
    "measurement_line = MeasurementLine([(0, 0), (1.8, 0)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot trajectories, geometry, and measurement areas/lines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from pedpy import plot_measurement_setup\n",
    "\n",
    "fig, axs = plt.subplots(3, int(len(trajectories) / 3), figsize=(10, 30))\n",
    "\n",
    "for (name, trajectory), ax in zip(\n",
    "    trajectories.items(), axs.ravel(), strict=False\n",
    "):\n",
    "    ax.set_title(name[:-4])\n",
    "\n",
    "    ax = plot_measurement_setup(\n",
    "        traj=trajectory,\n",
    "        walkable_area=walkable_area,\n",
    "        measurement_areas=[measurement_area],\n",
    "        measurement_lines=[measurement_line],\n",
    "        axes=ax,\n",
    "        traj_width=0.2,\n",
    "        traj_start_marker=\".\",\n",
    "        traj_end_marker=\"x\",\n",
    "        ma_color=\"g\",\n",
    "        ma_line_color=\"g\",\n",
    "        ma_alpha=0.2,\n",
    "        ml_color=\"b\",\n",
    "    )\n",
    "    ax.set_aspect(\"equal\")\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy.column_identifier import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Method A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute n-t and flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy import (\n",
    "    SpeedCalculation,\n",
    "    compute_flow,\n",
    "    compute_individual_speed,\n",
    "    compute_n_t,\n",
    ")\n",
    "\n",
    "nts = {}\n",
    "flows = {}\n",
    "\n",
    "for name, trajectory in trajectories.items():\n",
    "    individual_speed = compute_individual_speed(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        frame_step=10,\n",
    "        speed_calculation=SpeedCalculation.BORDER_SINGLE_SIDED,\n",
    "    )\n",
    "\n",
    "    nt, crossing = compute_n_t(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        measurement_line=measurement_line,\n",
    "    )\n",
    "\n",
    "    delta_frame = int(10 * trajectory.frame_rate)\n",
    "    flow = compute_flow(\n",
    "        nt=nt,\n",
    "        crossing_frames=crossing,\n",
    "        individual_speed=individual_speed,\n",
    "        delta_frame=delta_frame,\n",
    "        frame_rate=trajectory.frame_rate,\n",
    "    )\n",
    "\n",
    "    flows[name] = flow\n",
    "    nts[name] = nt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot n-t diagram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig = plt.figure(figsize=(7, 7))\n",
    "ax1 = fig.add_subplot(111)\n",
    "\n",
    "for name, nt in nts.items():\n",
    "    ax1.plot(nt.index, nt.cumulative_pedestrians, label=name)\n",
    "\n",
    "ax1.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot fundamental diagram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import image as mpimg\n",
    "\n",
    "style_options = {\n",
    "    \"uo-050-180-180\": {\"color\": \"red\", \"marker\": \"+\"},\n",
    "    \"uo-060-180-180\": {\"color\": \"limegreen\", \"marker\": \"x\"},\n",
    "    \"uo-070-180-180\": {\"color\": \"blue\", \"marker\": \"x\"},\n",
    "    \"uo-100-180-180\": {\n",
    "        \"color\": \"white\",\n",
    "        \"marker\": \"s\",\n",
    "        \"edgecolors\": \"hotpink\",\n",
    "    },\n",
    "    \"uo-145-180-180\": {\"color\": \"cyan\", \"marker\": \"s\"},\n",
    "    \"uo-180-180-070\": {\"color\": \"darkgrey\", \"marker\": \"^\"},\n",
    "    \"uo-180-180-095\": {\n",
    "        \"color\": \"white\",\n",
    "        \"marker\": \"^\",\n",
    "        \"edgecolors\": \"darkorange\",\n",
    "    },\n",
    "    \"uo-180-180-120\": {\n",
    "        \"color\": \"black\",\n",
    "        \"marker\": \"o\",\n",
    "    },\n",
    "    \"uo-180-180-180\": {\n",
    "        \"color\": \"white\",\n",
    "        \"marker\": \"o\",\n",
    "        \"edgecolors\": \"mediumorchid\",\n",
    "    },\n",
    "}\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2, figsize=(15, 6))\n",
    "\n",
    "for name, flow in flows.items():\n",
    "    ax0.scatter(\n",
    "        flow[FLOW_COL] / (flow[MEAN_SPEED_COL] * measurement_line.length),\n",
    "        flow[MEAN_SPEED_COL],\n",
    "        label=name,\n",
    "        **style_options[name],\n",
    "    )\n",
    "\n",
    "ax0.set_xlim(left=0, right=4)\n",
    "ax0.set_ylim(bottom=0, top=2.5)\n",
    "\n",
    "ax0.set_xlabel(\"rho / 1/m^2\")\n",
    "ax0.set_ylabel(\"v/ m/s\")\n",
    "ax0.grid()\n",
    "ax0.legend()\n",
    "\n",
    "img = mpimg.imread(folder / \"comparison/method_a_uo.png\")\n",
    "ax1.set_title(\"paper\")\n",
    "ax1.imshow(img)\n",
    "ax1.axis(\"off\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Method B"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute individual speed and density while passing the measurement area"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy import (\n",
    "    compute_classic_density,\n",
    "    compute_frame_range_in_area,\n",
    "    compute_passing_density,\n",
    "    compute_passing_speed,\n",
    ")\n",
    "\n",
    "passing_densities = {}\n",
    "passing_speeds = {}\n",
    "\n",
    "for name, trajectory in trajectories.items():\n",
    "    frames_in_area, passing_area = compute_frame_range_in_area(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        measurement_line=measurement_line,\n",
    "        width=2,\n",
    "    )\n",
    "\n",
    "    passing_speed = compute_passing_speed(\n",
    "        frames_in_area=frames_in_area,\n",
    "        frame_rate=trajectory.frame_rate,\n",
    "        distance=2.0,\n",
    "    )\n",
    "\n",
    "    density_per_frame = compute_classic_density(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        measurement_area=passing_area,\n",
    "    )\n",
    "    passing_density = compute_passing_density(\n",
    "        density_per_frame=density_per_frame, frames=frames_in_area\n",
    "    )\n",
    "\n",
    "    passing_speeds[name] = passing_speed\n",
    "    passing_densities[name] = passing_density"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot fundamental diagram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import image as mpimg\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2, figsize=(15, 6))\n",
    "\n",
    "for name in passing_speeds.keys():\n",
    "    ax0.scatter(\n",
    "        passing_densities[name][DENSITY_COL],\n",
    "        passing_speeds[name][SPEED_COL],\n",
    "        label=name,\n",
    "        **style_options[name],\n",
    "    )\n",
    "ax0.set_xlim(left=0, right=4)\n",
    "ax0.set_ylim(bottom=0, top=2.5)\n",
    "\n",
    "ax0.set_xlabel(\"rho / 1/m^2\")\n",
    "ax0.set_ylabel(\"v/ m/s\")\n",
    "ax0.grid()\n",
    "ax0.legend()\n",
    "\n",
    "img = mpimg.imread(folder / \"comparison/method_b_uo.png\")\n",
    "ax1.set_title(\"paper\")\n",
    "ax1.imshow(img)\n",
    "ax1.axis(\"off\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Method C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute density and speed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy import (\n",
    "    SpeedCalculation,\n",
    "    compute_classic_density,\n",
    "    compute_mean_speed_per_frame,\n",
    ")\n",
    "\n",
    "classic_densities = {}\n",
    "mean_area_speeds = {}\n",
    "\n",
    "for name, trajectory in trajectories.items():\n",
    "    individual_speed = compute_individual_speed(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        frame_step=5,\n",
    "        speed_calculation=SpeedCalculation.BORDER_SINGLE_SIDED,\n",
    "    )\n",
    "    mean_area_speed = compute_mean_speed_per_frame(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        measurement_area=measurement_area,\n",
    "        individual_speed=individual_speed,\n",
    "    )\n",
    "\n",
    "    classic_density = compute_classic_density(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        measurement_area=measurement_area,\n",
    "    )\n",
    "\n",
    "    classic_densities[name] = classic_density\n",
    "    mean_area_speeds[name] = mean_area_speed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot time series data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(\n",
    "    nrows=len(trajectories.values()), ncols=2, figsize=(20, 60)\n",
    ")\n",
    "row = 0\n",
    "\n",
    "ax[row, 1].set_title(\"Velocity\", size=\"xx-large\")\n",
    "ax[row, 0].set_title(\"Density\", size=\"xx-large\")\n",
    "\n",
    "for name, trajectory in trajectories.items():\n",
    "    ax[row, 0].annotate(\n",
    "        name,\n",
    "        xy=(0.5, 1),\n",
    "        xytext=(-ax[row, 0].yaxis.labelpad - 5, 0),\n",
    "        xycoords=ax[row, 0].yaxis.label,\n",
    "        textcoords=\"offset points\",\n",
    "        size=\"xx-large\",\n",
    "        ha=\"right\",\n",
    "        va=\"center\",\n",
    "        rotation=90,\n",
    "    )\n",
    "\n",
    "    ax[row, 0].plot(\n",
    "        classic_densities[name].index,\n",
    "        classic_densities[name][DENSITY_COL],\n",
    "        alpha=1,\n",
    "    )\n",
    "    ax[row, 0].set_xlim(left=0)\n",
    "    ax[row, 0].set_ylim(bottom=0, top=4)\n",
    "    ax[row, 0].set_xlabel(\"frame\")\n",
    "    ax[row, 0].set_ylabel(\"rho / 1/m^2\")\n",
    "    ax[row, 0].grid()\n",
    "\n",
    "    ax[row, 1].plot(\n",
    "        mean_area_speeds[name].index, mean_area_speeds[name], alpha=1\n",
    "    )\n",
    "    ax[row, 1].set_xlim(\n",
    "        left=0,\n",
    "    )\n",
    "    ax[row, 1].set_ylim(bottom=0, top=3)\n",
    "    ax[row, 1].set_xlabel(\"frame\")\n",
    "    ax[row, 1].set_ylabel(\"v / m/s\")\n",
    "    ax[row, 1].grid()\n",
    "\n",
    "    row += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot fundamental diagram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import image as mpimg\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2, figsize=(15, 6))\n",
    "\n",
    "ax0.set_title(\"Fundamental Diagram\")\n",
    "\n",
    "for name in classic_densities.keys():\n",
    "    ax0.scatter(\n",
    "        classic_densities[name][DENSITY_COL],\n",
    "        mean_area_speeds[name],\n",
    "        alpha=1,\n",
    "        label=name,\n",
    "        **style_options[name],\n",
    "    )\n",
    "ax0.set_xlim(left=0, right=4)\n",
    "ax0.set_ylim(bottom=0, top=2.5)\n",
    "\n",
    "ax0.set_xlabel(\"rho / 1/m^2\")\n",
    "ax0.set_ylabel(\"v/ m/s\")\n",
    "ax0.grid()\n",
    "ax0.legend()\n",
    "\n",
    "img = mpimg.imread(folder / \"comparison/method_c_uo.png\")\n",
    "ax1.set_title(\"paper\")\n",
    "ax1.imshow(img)\n",
    "ax1.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Method D"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute density and speed"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Without cutoff radius"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy import (\n",
    "    SpeedCalculation,\n",
    "    compute_individual_voronoi_polygons,\n",
    "    compute_voronoi_density,\n",
    "    compute_voronoi_speed,\n",
    ")\n",
    "\n",
    "voronoi_densities = {}\n",
    "voronoi_speeds = {}\n",
    "individual_speeds = {}\n",
    "individuals = {}\n",
    "\n",
    "for name, trajectory in trajectories.items():\n",
    "    individual = compute_individual_voronoi_polygons(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        walkable_area=walkable_area,\n",
    "    )\n",
    "    voronoi_density, intersecting = compute_voronoi_density(\n",
    "        individual_voronoi_data=individual,\n",
    "        measurement_area=measurement_area,\n",
    "    )\n",
    "\n",
    "    individual_speed = compute_individual_speed(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        frame_step=5,\n",
    "        speed_calculation=SpeedCalculation.BORDER_SINGLE_SIDED,\n",
    "    )\n",
    "    voronoi_speed = compute_voronoi_speed(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        individual_voronoi_intersection=intersecting,\n",
    "        measurement_area=measurement_area,\n",
    "        individual_speed=individual_speed,\n",
    "    )\n",
    "\n",
    "    voronoi_densities[name] = voronoi_density\n",
    "    voronoi_speeds[name] = voronoi_speed\n",
    "    individual_speeds[name] = individual_speed\n",
    "    individuals[name] = intersecting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### With cutoff radius"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy import (\n",
    "    Cutoff,\n",
    "    SpeedCalculation,\n",
    "    compute_voronoi_density,\n",
    "    compute_voronoi_speed,\n",
    ")\n",
    "\n",
    "voronoi_densities_cutoff = {}\n",
    "voronoi_speeds_cutoff = {}\n",
    "individual_speeds_cutoff = {}\n",
    "individual_cutoffs = {}\n",
    "for name, trajectory in trajectories.items():\n",
    "    individual_cutoff = compute_individual_voronoi_polygons(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        walkable_area=walkable_area,\n",
    "        cut_off=Cutoff(radius=0.8, quad_segments=3),\n",
    "    )\n",
    "    voronoi_density_cutoff, intersecting_cutoff = compute_voronoi_density(\n",
    "        individual_voronoi_data=individual_cutoff,\n",
    "        measurement_area=measurement_area,\n",
    "    )\n",
    "\n",
    "    individual_speed = compute_individual_speed(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        frame_step=5,\n",
    "        speed_calculation=SpeedCalculation.BORDER_SINGLE_SIDED,\n",
    "    )\n",
    "\n",
    "    voronoi_speed_cutoff = compute_voronoi_speed(\n",
    "        traj_data=trajectories_in_steady_state[name],\n",
    "        individual_voronoi_intersection=intersecting_cutoff,\n",
    "        individual_speed=individual_speed,\n",
    "        measurement_area=measurement_area,\n",
    "    )\n",
    "\n",
    "    voronoi_densities_cutoff[name] = voronoi_density_cutoff\n",
    "    voronoi_speeds_cutoff[name] = voronoi_speed_cutoff\n",
    "    individual_speeds_cutoff[name] = individual_speed\n",
    "    individual_cutoffs[name] = intersecting_cutoff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot time-series data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(\n",
    "    nrows=len(trajectories.values()), ncols=2, figsize=(20, 60)\n",
    ")\n",
    "row = 0\n",
    "\n",
    "ax[row, 1].set_title(\"Speed\", size=\"xx-large\")\n",
    "ax[row, 0].set_title(\"Density\", size=\"xx-large\")\n",
    "\n",
    "for name, trajectory in trajectories.items():\n",
    "    ax[row, 0].annotate(\n",
    "        name,\n",
    "        xy=(0.5, 1),\n",
    "        xytext=(-ax[row, 0].yaxis.labelpad - 5, 0),\n",
    "        xycoords=ax[row, 0].yaxis.label,\n",
    "        textcoords=\"offset points\",\n",
    "        size=\"xx-large\",\n",
    "        ha=\"right\",\n",
    "        va=\"center\",\n",
    "        rotation=90,\n",
    "    )\n",
    "\n",
    "    ax[row, 0].plot(\n",
    "        voronoi_densities[name].index,\n",
    "        voronoi_densities[name][DENSITY_COL],\n",
    "        alpha=1,\n",
    "        label=\"without cut-off\",\n",
    "    )\n",
    "    ax[row, 0].plot(\n",
    "        voronoi_densities_cutoff[name].index,\n",
    "        voronoi_densities_cutoff[name][DENSITY_COL],\n",
    "        alpha=1,\n",
    "        label=\"with cut-off\",\n",
    "    )\n",
    "    ax[row, 0].set_xlim(left=0)\n",
    "    ax[row, 0].set_ylim(bottom=0, top=4)\n",
    "    ax[row, 0].set_xlabel(\"frame\")\n",
    "    ax[row, 0].set_ylabel(\"rho / 1/m^2\")\n",
    "    ax[row, 0].grid()\n",
    "    ax[row, 0].legend()\n",
    "\n",
    "    ax[row, 1].plot(\n",
    "        voronoi_speeds[name].index,\n",
    "        voronoi_speeds[name],\n",
    "        alpha=1,\n",
    "        label=\"without cut-off\",\n",
    "    )\n",
    "    ax[row, 1].plot(\n",
    "        voronoi_speeds_cutoff[name].index,\n",
    "        voronoi_speeds_cutoff[name],\n",
    "        alpha=1,\n",
    "        label=\"with cut-off\",\n",
    "    )\n",
    "    ax[row, 1].set_xlim(\n",
    "        left=0,\n",
    "    )\n",
    "    ax[row, 1].set_ylim(bottom=0, top=3)\n",
    "    ax[row, 1].set_xlabel(\"frame\")\n",
    "    ax[row, 1].set_ylabel(\"v / m/s\")\n",
    "    ax[row, 1].grid()\n",
    "    ax[row, 1].legend()\n",
    "    row += 1\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot fundamental diagram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import image as mpimg\n",
    "\n",
    "fig, (ax0, ax1, ax2) = plt.subplots(nrows=1, ncols=3, figsize=(20, 6))\n",
    "fig.suptitle(\"Fundamental diagram\")\n",
    "\n",
    "ax0.set_title(\"without cutoff\")\n",
    "for name in voronoi_densities.keys():\n",
    "    ax0.scatter(\n",
    "        voronoi_densities[name][DENSITY_COL],\n",
    "        voronoi_speeds[name],\n",
    "        alpha=1,\n",
    "        label=name,\n",
    "        **style_options[name],\n",
    "    )\n",
    "ax0.set_xlim(left=0, right=4)\n",
    "ax0.set_ylim(bottom=0, top=2.5)\n",
    "\n",
    "ax0.set_xlabel(\"rho / 1/m^2\")\n",
    "ax0.set_ylabel(\"v/ m/s\")\n",
    "ax0.grid()\n",
    "ax0.legend()\n",
    "\n",
    "ax1.set_title(\"with cutoff\")\n",
    "for name in voronoi_densities_cutoff.keys():\n",
    "    ax1.scatter(\n",
    "        voronoi_densities_cutoff[name][DENSITY_COL],\n",
    "        voronoi_speeds_cutoff[name],\n",
    "        alpha=1,\n",
    "        label=name,\n",
    "        **style_options[name],\n",
    "    )\n",
    "ax1.set_xlim(left=0, right=4)\n",
    "ax1.set_ylim(bottom=0, top=2.5)\n",
    "\n",
    "ax1.set_xlabel(\"rho / 1/m^2\")\n",
    "ax1.set_ylabel(\"v/ m/s\")\n",
    "ax1.grid()\n",
    "ax1.legend()\n",
    "\n",
    "img = mpimg.imread(folder / \"comparison/method_d_uo.png\")\n",
    "ax2.set_title(\"paper\")\n",
    "ax2.imshow(img)\n",
    "ax2.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute speed and density profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from pedpy import (\n",
    "    SpeedCalculation,\n",
    "    SpeedMethod,\n",
    "    compute_individual_voronoi_polygons,\n",
    "    compute_profiles,\n",
    "    compute_voronoi_density,\n",
    "    compute_voronoi_speed,\n",
    ")\n",
    "\n",
    "trajectory = list(trajectories.values())[7]\n",
    "\n",
    "frames_data = TrajectoryData(\n",
    "    data=trajectory.data[trajectory.data.frame.between(750, 1500)],\n",
    "    frame_rate=trajectory.frame_rate,\n",
    ")\n",
    "individual_voronoi = compute_individual_voronoi_polygons(\n",
    "    traj_data=frames_data,\n",
    "    walkable_area=walkable_area,\n",
    "    cut_off=Cutoff(radius=0.89, quad_segments=3),\n",
    ")\n",
    "voronoi_density_frames, individual_frames = compute_voronoi_density(\n",
    "    individual_voronoi_data=individual_voronoi,\n",
    "    measurement_area=measurement_area,\n",
    ")\n",
    "\n",
    "individual_speed = compute_individual_speed(\n",
    "    traj_data=trajectory,\n",
    "    frame_step=5,\n",
    "    speed_calculation=SpeedCalculation.BORDER_SINGLE_SIDED,\n",
    ")\n",
    "voronoi_speed_frames = compute_voronoi_speed(\n",
    "    traj_data=frames_data,\n",
    "    individual_voronoi_intersection=individual_frames,\n",
    "    individual_speed=individual_speed[\n",
    "        individual_speed.frame.between(750, 1500)\n",
    "    ],\n",
    "    measurement_area=measurement_area,\n",
    ")\n",
    "\n",
    "density_profiles, speed_profiles = compute_profiles(\n",
    "    data=pd.merge(\n",
    "        individual_frames,\n",
    "        individual_speed,\n",
    "        on=[ID_COL, FRAME_COL],\n",
    "        how=\"left\",\n",
    "    ),\n",
    "    walkable_area=walkable_area,\n",
    "    grid_size=0.25,\n",
    "    speed_method=SpeedMethod.VORONOI,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "bounds = walkable_area.bounds\n",
    "\n",
    "fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2, figsize=(10, 10))\n",
    "\n",
    "ax0.set_title(\"Density\")\n",
    "cm = ax0.imshow(\n",
    "    np.mean(density_profiles, axis=0),\n",
    "    extent=[bounds[0], bounds[2], bounds[1], bounds[3]],\n",
    "    interpolation=\"None\",\n",
    "    cmap=\"jet\",\n",
    "    vmin=0,\n",
    "    vmax=5,\n",
    ")\n",
    "fig.colorbar(cm, ax=ax0, shrink=0.3)\n",
    "ax0.plot(*walkable_area.polygon.exterior.xy, color=\"w\")\n",
    "\n",
    "ax1.set_title(\"Velocity\")\n",
    "cm = ax1.imshow(\n",
    "    np.mean(speed_profiles, axis=0),\n",
    "    extent=[bounds[0], bounds[2], bounds[1], bounds[3]],\n",
    "    interpolation=\"None\",\n",
    "    cmap=\"jet\",\n",
    "    vmin=0,\n",
    "    vmax=1.5,\n",
    ")\n",
    "fig.colorbar(cm, ax=ax1, shrink=0.3)\n",
    "\n",
    "ax1.plot(*walkable_area.polygon.exterior.xy, color=\"w\")\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Voronoi cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pedpy import plot_voronoi_cells\n",
    "from pedpy.column_identifier import SPEED_COL\n",
    "\n",
    "# Choose trajectory file to plot\n",
    "name = list(trajectories.keys())[7]\n",
    "\n",
    "# Choose frame to plot\n",
    "frame = 700\n",
    "\n",
    "# Choose which voronoi data to plot\n",
    "use_cut_off_data = True\n",
    "\n",
    "# merge all relevant data into one data frame\n",
    "voronoi_data = (\n",
    "    individual_cutoffs[name] if use_cut_off_data else individuals[name]\n",
    ")\n",
    "speed_data = (\n",
    "    individual_speeds_cutoff[name][[ID_COL, FRAME_COL, SPEED_COL]]\n",
    "    if use_cut_off_data\n",
    "    else individual_speeds[name][[ID_COL, FRAME_COL, SPEED_COL]]\n",
    ")\n",
    "combined_data = pd.merge(speed_data, voronoi_data, on=[ID_COL, FRAME_COL])\n",
    "\n",
    "# Start plotting\n",
    "fig = plt.figure(layout=\"constrained\")\n",
    "ax1 = fig.add_subplot(131, aspect=\"equal\")\n",
    "ax1.set_title(\"Color by density\")\n",
    "plot_voronoi_cells(\n",
    "    voronoi_data=combined_data,\n",
    "    walkable_area=walkable_area,\n",
    "    measurement_area=measurement_area,\n",
    "    frame=frame,\n",
    "    traj_data=trajectories[name],\n",
    "    color_by_column=DENSITY_COL,\n",
    "    ped_size=5,\n",
    "    voronoi_outside_ma_alpha=0.4,\n",
    "    axes=ax1,\n",
    ")\n",
    "\n",
    "ax2 = fig.add_subplot(132, aspect=\"equal\")\n",
    "ax2.set_title(\"Color by speed\")\n",
    "plot_voronoi_cells(\n",
    "    voronoi_data=combined_data,\n",
    "    walkable_area=walkable_area,\n",
    "    measurement_area=measurement_area,\n",
    "    frame=frame,\n",
    "    traj_data=trajectories[name],\n",
    "    color_by_column=SPEED_COL,\n",
    "    ped_size=5,\n",
    "    voronoi_outside_ma_alpha=0.4,\n",
    "    axes=ax2,\n",
    ")\n",
    "\n",
    "ax3 = fig.add_subplot(133, aspect=\"equal\")\n",
    "ax3.set_title(\"Color by id\")\n",
    "plot_voronoi_cells(\n",
    "    voronoi_data=combined_data,\n",
    "    walkable_area=walkable_area,\n",
    "    measurement_area=measurement_area,\n",
    "    frame=frame,\n",
    "    traj_data=trajectories[name],\n",
    "    color_by_column=ID_COL,\n",
    "    ped_size=5,\n",
    "    voronoi_outside_ma_alpha=0.4,\n",
    "    axes=ax3,\n",
    ")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}